diff --git a/.github/workflows/pr-test.yml b/.github/workflows/pr-test.yml index efbfe8636..bc8b8e769 100644 --- a/.github/workflows/pr-test.yml +++ b/.github/workflows/pr-test.yml @@ -224,7 +224,7 @@ jobs: strategy: fail-fast: false matrix: - info: [{"num_gpus": 8, "test_file": "test_quick_start_glm4_9B.py"}, {"num_gpus": 8, "test_file": "test_glm4.7_30B_A3B_pd_mooncake.py"}, {"num_gpus": 8, "test_file": "test_qwen3_30B_A3B.py", "use_deepep": "1", "use_fp8_rollout": "1"}, {"num_gpus": 8, "test_file": "test_qwen3.6_35B_A3B_pd_mooncake.py", "use_deepep": "1"}, {"enable_eval": "0", "num_gpus": 8, "test_file": "test_qwen3_30B_A3B_r3.py", "use_deepep": "1", "use_fp8_rollout": "1"}, {"enable_eval": "0", "num_gpus": 8, "test_file": "test_qwen3_30B_A3B_r3.py"}, {"num_gpus": 8, "test_file": "test_qwen3_4B_ppo.py"}, {"num_gpus": 8, "test_file": "test_qwen3_4B_ppo_disaggregate.py"}, {"num_gpus": 8, "test_file": "test_qwen3_4B_ppo_train_critic_only.py"}, {"num_gpus": 8, "test_file": "test_moonlight_16B_A3B.py"}, {"enable_eval": "0", "num_gpus": 8, "test_file": "test_moonlight_16B_A3B_r3.py"}, {"num_gpus": 8, "test_file": "test_mimo_7B_mtp_only_grad.py"}, {"num_gpus": 8, "test_file": "test_qwen2.5_0.5B_debug_rollout_then_train.py"}, {"num_gpus": 8, "test_file": "test_qwen2.5_0.5B_opd_vllm.py"}] + info: [{"num_gpus": 8, "test_file": "test_quick_start_glm4_9B.py"}, {"num_gpus": 8, "test_file": "test_glm4.7_30B_A3B_pd_mooncake.py"}, {"num_gpus": 8, "test_file": "test_qwen3_30B_A3B.py", "use_deepep": "1", "use_fp8_rollout": "1"}, {"num_gpus": 8, "test_file": "test_qwen3.6_35B_A3B_pd_mooncake.py", "use_deepep": "1"}, {"enable_eval": "0", "num_gpus": 8, "test_file": "test_qwen3_30B_A3B_r3.py", "use_deepep": "1", "use_fp8_rollout": "1"}, {"enable_eval": "0", "num_gpus": 8, "test_file": "test_qwen3_30B_A3B_r3.py"}, {"num_gpus": 8, "test_file": "test_qwen3_4B_ppo.py"}, {"num_gpus": 8, "test_file": "test_qwen3_4B_ppo_disaggregate.py"}, {"num_gpus": 8, "test_file": "test_qwen3_4B_ppo_train_critic_only.py"}, {"num_gpus": 8, "test_file": "test_moonlight_16B_A3B.py"}, {"enable_eval": "0", "num_gpus": 8, "test_file": "test_moonlight_16B_A3B_r3.py"}, {"num_gpus": 8, "test_file": "test_qwen2.5_0.5B_debug_rollout_then_train.py"}, {"num_gpus": 8, "test_file": "test_qwen2.5_0.5B_opd_vllm.py"}] defaults: run: working-directory: ${{ github.workspace }} @@ -542,7 +542,7 @@ jobs: strategy: fail-fast: false matrix: - info: [{"num_gpus": 4, "test_file": "test_qwen3.5_0.8B_gsm8k_async_short.py"}, {"num_gpus": 4, "test_file": "test_qwen3.5_0.8B_gsm8k_short.py"}, {"num_gpus": 8, "test_file": "test_quick_start_glm4_9B.py"}, {"num_gpus": 8, "test_file": "test_glm4.7_30B_A3B_pd_mooncake.py"}, {"num_gpus": 8, "test_file": "test_qwen3_30B_A3B.py"}, {"num_gpus": 8, "test_file": "test_qwen3.6_35B_A3B_pd_mooncake.py", "use_deepep": "1"}, {"num_gpus": 8, "test_file": "test_qwen3_4B_ppo.py"}, {"num_gpus": 8, "test_file": "test_moonlight_16B_A3B.py"}, {"num_gpus": 8, "test_file": "test_mimo_7B_mtp_only_grad.py"}, {"num_gpus": 8, "test_file": "test_qwen3_0.6B_parallel_check.py"}, {"num_gpus": 8, "test_file": "test_qwen3_4B_ckpt.py"}, {"num_gpus": 8, "test_args": "--async-save", "test_file": "test_qwen3_4B_ckpt.py"}, {"num_gpus": 8, "test_file": "test_qwen2.5_0.5B_debug_rollout_then_train.py"}, {"num_gpus": 8, "test_file": "test_qwen2.5_0.5B_opd_vllm.py"}] + info: [{"num_gpus": 4, "test_file": "test_qwen3.5_0.8B_gsm8k_async_short.py"}, {"num_gpus": 4, "test_file": "test_qwen3.5_0.8B_gsm8k_short.py"}, {"num_gpus": 8, "test_file": "test_quick_start_glm4_9B.py"}, {"num_gpus": 8, "test_file": "test_glm4.7_30B_A3B_pd_mooncake.py"}, {"num_gpus": 8, "test_file": "test_qwen3_30B_A3B.py"}, {"num_gpus": 8, "test_file": "test_qwen3.6_35B_A3B_pd_mooncake.py", "use_deepep": "1"}, {"num_gpus": 8, "test_file": "test_qwen3_4B_ppo.py"}, {"num_gpus": 8, "test_file": "test_moonlight_16B_A3B.py"}, {"num_gpus": 8, "test_file": "test_qwen3_0.6B_parallel_check.py"}, {"num_gpus": 8, "test_file": "test_qwen3_4B_ckpt.py"}, {"num_gpus": 8, "test_args": "--async-save", "test_file": "test_qwen3_4B_ckpt.py"}, {"num_gpus": 8, "test_file": "test_qwen2.5_0.5B_debug_rollout_then_train.py"}, {"num_gpus": 8, "test_file": "test_qwen2.5_0.5B_opd_vllm.py"}] defaults: run: working-directory: ${{ github.workspace }} @@ -767,4 +767,4 @@ jobs: else python tests/ci/gpu_lock_exec.py --count "$NUM_GPUS" -- python "$TEST_PATH" "${TEST_ARGS_ARRAY[@]}" fi - ' \ No newline at end of file + ' diff --git a/.github/workflows/pr-test.yml.j2 b/.github/workflows/pr-test.yml.j2 index 2f0f34d95..adab45a63 100644 --- a/.github/workflows/pr-test.yml.j2 +++ b/.github/workflows/pr-test.yml.j2 @@ -31,7 +31,6 @@ {'test_file': 'test_qwen3_4B_ppo_train_critic_only.py', 'num_gpus': 8}, {'test_file': 'test_moonlight_16B_A3B.py', 'num_gpus': 8}, {'test_file': 'test_moonlight_16B_A3B_r3.py', 'num_gpus': 8, 'enable_eval': '0'}, - {'test_file': 'test_mimo_7B_mtp_only_grad.py', 'num_gpus': 8}, {'test_file': 'test_qwen2.5_0.5B_debug_rollout_then_train.py', 'num_gpus': 8}, {'test_file': 'test_qwen2.5_0.5B_opd_vllm.py', 'num_gpus': 8}, ], @@ -75,7 +74,6 @@ {'test_file': 'test_qwen3.6_35B_A3B_pd_mooncake.py', 'num_gpus': 8, 'use_deepep': '1'}, {'test_file': 'test_qwen3_4B_ppo.py', 'num_gpus': 8}, {'test_file': 'test_moonlight_16B_A3B.py', 'num_gpus': 8}, - {'test_file': 'test_mimo_7B_mtp_only_grad.py', 'num_gpus': 8}, {'test_file': 'test_qwen3_0.6B_parallel_check.py', 'num_gpus': 8}, {'test_file': 'test_qwen3_4B_ckpt.py', 'num_gpus': 8}, {'test_file': 'test_qwen3_4B_ckpt.py', 'test_args': '--async-save', 'num_gpus': 8}, diff --git a/docs/en/advanced/low-precision.md b/docs/en/advanced/low-precision.md deleted file mode 100644 index 951d5c163..000000000 --- a/docs/en/advanced/low-precision.md +++ /dev/null @@ -1,133 +0,0 @@ -# Low Precision Training - -- [FP8 rollout and FP8 training](#FP8-rollout-and-BF16-training) -- [FP8 rollout and FP8 training](#FP8-rollout-and-FP8-training) -- [INT4 QAT Training](#INT4-QAT-Training) - -## FP8 rollout and BF16 training - -You can run FP8 rollout simply by setting `--hf-checkpoint` with an blockwise quantized huggingface checkpoint, which can be converted by: - -```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 -``` - -Please ensure that the converted checkpoint points to a directory where the `config.json` contains the correct `quantization_config` so that slime can automatically use FP8 quantization during weight updates. - -## FP8 rollout and FP8 training - -We also observed that under FP8 training and inference, it can achieve more efficient inference throughput and lower training-inference mismatch, resulting in more stable training. More details can be found in [this blog](https://lmsys.org/blog/2025-11-25-fp8-rl/). - -### Quick Start - -1. Convert your HuggingFace model weights to FP8 format using the above `tools/convert_hf_to_fp8.py`. - -2. Setting up the running script: - -For training tasks, we need to add these flags: - -```bash ---fp8-format e4m3 ---fp8-recipe blockwise -# --fp8-param-gather # [optional] Currently incompatible with CPU Adam -``` - -Then ensure the `NVTE_FP8_BLOCK_SCALING_FP32_SCALES` environment variable is enabled. - -Note that only `Linear` and `GroupLinear` layers in TransformerEngine use fp8 format. `embedding` and `lm_head` remain in their original precision. If `--fp8-param-gather` is not enabled, weights in TransformerEngine remain in bf16 format, only being cast to fp8 format during `GEMM` or `GroupGEMM` operations. - -3. Start FP8 training with - -```bash -# Qwen3-4B Int4 training -bash scripts/low_precision/run-qwen3-4b-fp8.sh - -# Qwen3-30B-A3B (2 nodes) -bash scripts/low_precision/run-qwen3-30b-a3b-fp8.sh -``` - -4. Use the saved checkpoint for evaluation. - -Note that TransformerEngine does not specifically save FP8 quantized weights; the saved torch dist remains in original precision (usually bf16). If you want to evaluate under FP8, you need to convert the checkpoint from `torch_dist` to HuggingFace format, then convert to FP8 HuggingFace format. - - -### Quick Explanation - -Here's a quick explanation of how FP8 training is currently implemented in slime: - -1. Initialization: If FP8 recipe is enabled, layers will be built in FP8 context. - -2. Training: During training, weights and activations are quantized online to nvfp8 format, and cuBLAS FP8 GEMM is called for various GEMM computations in forward and backward passes. - -3. Weight updates: During RL weight updates, Megatron first dequantizes FP8 weights to bf16 format, then slime quantizes these bf16 weights to fp8 format and sends them to vLLM. (This additional dequantization and quantization is not elegant, but we haven't modified the interface yet for framework compatibility.) - -4. Save checkpoint: Similar to weight updates, if checkpoints need to be saved from the training engine, they will also be dequantized back to bf16 and saved to `torch_dist` format checkpoints. - - -### TODO - -Currently, FP8 is far from being a complete feature and still has the following bugs, for examples: - -- FP8 weights (`--fp8-param-gather`) can provide memory savings benefits, but currently FP8 weights must be used with TransformerEngine's FusedAdam, which conflicts with the commonly used Adam CPU offload technique in Megatron-LM. - -## INT4 QAT Training - -This guide provides examples for INT4 STE (Straight-Through Estimator) training and INT4 inference. Utilizing INT4 inference significantly improves throughput, thereby accelerating the training pipeline (specifically during the rollout generation phase). - -### Quick Start - -1. Convert HuggingFace Weights to INT4 -Use the `tools/convert_hf_to_int4_direct.py` script to convert BF16 weights to INT4 format. Ensure that the `--hf-checkpoint` parameter points to a directory where `config.json` contains the correct `quantization_config`. slime will automatically utilize INT4 quantization during weight updates. - -```bash -python tools/convert_hf_to_int4_direct.py \ - --model-dir /path/to/your/original/models \ - --save-dir /path/to/your/save/models \ -``` - -Note: If you only hope to run with INT4 rollout, you only need to set the `--hf-checkpoint` to the converted INT4 checkpoint. - -2. Start INT4 QAT Training - -You need to configure the specific environment variables for quantization settings. - -**Environment Variables:** - -* **`OPEN_TRAINING_INT4_FAKE_QAT_FLAG`**: Enables fake quantization operations for INT4 training. -* **`OPEN_TRAINING_INT4_GROUP_SIZE`**: Specifies the block size (group size) for model quantization. - * Set to **128** for `moonlight-16B-A3B` 、 `qwen3-30B-A3B`and `qwen3-235B-A22B-int4`. - * Set to **32** for `kimi-k2-Thinking-int4`. - -**Configuration Example:** - -```json -RUNTIME_ENV_JSON="{ - \"env_vars\": { - ... - \"OPEN_TRAINING_INT4_FAKE_QAT_FLAG\": \"1\", - \"OPEN_TRAINING_INT4_GROUP_SIZE\": \"128\" - } -}" -``` - -**Launch Commands:** - -```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, please start the Ray service according to your cluster configuration. diff --git a/docs/en/advanced/on-policy-distillation.md b/docs/en/advanced/on-policy-distillation.md deleted file mode 100644 index b4efae4c0..000000000 --- a/docs/en/advanced/on-policy-distillation.md +++ /dev/null @@ -1,111 +0,0 @@ -# On-Policy Distillation - -On-policy distillation (OPD) enables a student model to learn from a larger teacher model by training on its own rollouts while matching the teacher's token-level log-probabilities. OPD is orthogonal to advantage estimators — it works as an additive KL penalty on top of any estimator (GRPO, PPO, REINFORCE++, etc.). - -## Key Arguments - -| Argument | Description | -|----------|-------------| -| `--use-opd` | Enable on-policy distillation. Required flag to use OPD. | -| `--opd-type` | Type of OPD: `vllm` or `megatron`. Required when `--use-opd` is set. | -| `--opd-kl-coef` | OPD KL penalty coefficient (default: 1.0). Controls the weight of the distillation signal relative to the RL advantage. | -| `--opd-teacher-load` | Path to teacher Megatron checkpoint. **Required** when `--opd-type=megatron`, **must not be set** when `--opd-type=vllm`. | -| `--opd-teacher-ckpt-step` | Optional checkpoint step for teacher model. | - -## How It Works - -OPD modifies the advantage computation by subtracting a KL penalty term that encourages the student to match the teacher's output distribution: - -$$ -\hat{A}_t = A_t - \lambda_{\text{opd}} \cdot D_{\text{KL}}(P_{\text{teacher}} \| P_{\text{student}})_t -$$ - -Where $A_t$ is the original advantage from the base estimator (e.g., GRPO), $\lambda_{\text{opd}}$ is `--opd-kl-coef`, and $D_{\text{KL}}$ is the token-level reverse KL divergence. - -This means OPD can be combined with any advantage estimator, including GRPO, PPO, REINFORCE++, and GSPO. - -## Two Teacher Modes - -### vLLM Mode (`--opd-type vllm`) - -The teacher runs on an external vLLM server. Teacher log-probs are obtained during the rollout phase. - -**When to use**: The teacher has a different architecture from the student, or the teacher is too large to load alongside the training model. - -**How it works**: -1. An external vLLM server runs the teacher model. -2. During rollout, the custom reward function (`slime.rollout.on_policy_distillation.reward_func`) sends each sample to the teacher server to obtain token-level log-probs. -3. The custom post-processing function (`slime.rollout.on_policy_distillation.post_process_rewards`) trims the teacher log-probs to the response span and stores them in `sample.teacher_log_probs`. -4. During training, the KL penalty is computed from the stored teacher log-probs and applied to advantages. - -**Configuration**: -```bash ---use-opd ---opd-type vllm ---opd-kl-coef 1.0 ---custom-rm-path slime.rollout.on_policy_distillation.reward_func ---custom-reward-post-process-path slime.rollout.on_policy_distillation.post_process_rewards ---rm-url http://:/inference/v1/generate -``` - -### Megatron Mode (`--opd-type megatron`) - -The teacher model is loaded directly into Megatron via `--opd-teacher-load`. Teacher log-probs are computed during the training forward pass. - -**When to use**: The teacher has the same architecture as the student/reference model and fits in GPU memory. - -**How it works**: -1. The teacher model is loaded as an additional Megatron model during initialization. -2. During the training forward pass, the teacher model computes log-probs for each sample. -3. The KL penalty is computed inline and applied to advantages. - -**Configuration**: -```bash ---use-opd ---opd-type megatron ---opd-kl-coef 1.0 ---opd-teacher-load /path/to/teacher_torch_dist -``` - -> **Note**: The teacher checkpoint must be in Megatron format (`torch_dist` or `torch`). You can convert from HuggingFace format using `tools/convert_hf_to_torch_dist.py`. - -## Running the Examples - -Complete example scripts are provided in `examples/on_policy_distillation/`: - -### vLLM Teacher - -```bash -# 1. Download models and data -hf download Qwen/Qwen3-32B --local-dir /root/Qwen3-32B -hf download Qwen/Qwen3-8B --local-dir /root/Qwen3-8B -hf download --repo-type dataset zhuzilin/dapo-math-17k --local-dir /root/dapo-math-17k - -# 2. Convert student model -cd /root/slime -source scripts/models/qwen3-8B.sh -PYTHONPATH=/root/Megatron-LM python tools/convert_hf_to_torch_dist.py \ - ${MODEL_ARGS[@]} \ - --hf-checkpoint /root/Qwen3-8B \ - --save /root/Qwen3-8B_torch_dist - -# 3. Run -bash examples/on_policy_distillation/run-qwen3-8B-opd.sh -``` - -### Megatron Teacher - -```bash -# 1. Convert both student and teacher models to Megatron format -# 2. Run -bash examples/on_policy_distillation/run-qwen3-8B-opd-megatron.sh -``` - -## Preliminary Results - -Using Qwen3-8B-Base model SFT-ed on part of the [OpenThoughts3-1.2M](https://huggingface.co/datasets/open-thoughts/OpenThoughts3-1.2M) dataset, on-policy distillation with a Qwen3-32B teacher on the remaining data yields: - -| | Pass@1 | -|-----------------------------------------------|--------| -| Qwen3-8B-Base + SFT | 76% | -| Qwen3-8B-Base + SFT + On-Policy Distillation | 94% | diff --git a/docs/en/developer_guide/ci.md b/docs/en/developer_guide/ci.md index 82ee9cbcc..700daf8ce 100644 --- a/docs/en/developer_guide/ci.md +++ b/docs/en/developer_guide/ci.md @@ -20,7 +20,7 @@ Add a label to your PR to trigger the corresponding test suite: | Label | Job | Description | |---|---|---| | `run-ci-short` | `e2e-test-short` | Lightweight smoke tests with Qwen2.5-0.5B (4 GPUs). Fast feedback loop. | -| `run-ci-megatron` | `e2e-test-megatron` | Core Megatron training tests covering dense, MoE, PPO, MTP, OPD, etc. | +| `run-ci-megatron` | `e2e-test-megatron` | Core Megatron training tests covering dense, MoE, PPO, OPD, etc. | | `run-ci-precision` | `e2e-test-precision` | Numerical precision validation (parallel check). | | `run-ci-ckpt` | `e2e-test-ckpt` | Checkpoint save/load correctness (sync and async-save). | | `run-ci-image` | `e2e-test-image` | Full test suite run on `inferactinc/public:vime-vllm-cu129-latest` image (for image validation). | @@ -57,7 +57,7 @@ This is the primary label for validating Megatron-backend changes. It covers: - Dense models: GLM4-9B, Qwen3-4B (PPO) - MoE models: Qwen3-30B-A3B (with DeepEP + FP8), Qwen3.6-35B-A3B PD + Mooncake, Moonlight-16B-A3B -- Specialized: MiMo-7B MTP, Qwen2.5-0.5B debug rollout-then-train, OPD with vLLM teacher +- Specialized: Qwen2.5-0.5B debug rollout-then-train, OPD with vLLM teacher All tests use 8 GPUs. If you are modifying Megatron training logic, loss computation, or checkpoint conversion, this is the label to use. diff --git a/docs/en/examples/deepseek-r1.md b/docs/en/examples/deepseek-r1.md deleted file mode 100644 index acbbf1135..000000000 --- a/docs/en/examples/deepseek-r1.md +++ /dev/null @@ -1,202 +0,0 @@ -# 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 with DP. 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 slime/ -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 slime/ -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 slime/ -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/THUDM/slime/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/THUDM/slime/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_slime/ - --save $BASE_DIR/DeepSeek-R1_slime/ - --save-interval 20 -) -``` - -slime 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 slime. - -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. - -⚠️ slime 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 slime: - -```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` corresponds to vLLM's `tp_size`. Other vLLM parameters are passed to slime by adding a `--vllm-` prefix. - -`--vllm-server-concurrency` is a parameter specific to slime. It is used to prevent the vLLM engine'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 - - # data parallelism for the attention block - --vllm-data-parallel-size 8 - - # make every dp rank have 128 concurrency - --vllm-server-concurrency 1024 -) -``` - -#### 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 -) -``` diff --git a/docs/en/examples/glm4-9B.md b/docs/en/examples/glm4-9B.md deleted file mode 100644 index 9ea7d3883..000000000 --- a/docs/en/examples/glm4-9B.md +++ /dev/null @@ -1,280 +0,0 @@ -# GLM4-9B with 8xH100 - -## Environment Setup - -After pulling the `inferactinc/public:vime-vllm-cu129-latest` image, initialize the image environment as follows: - -```bash -cd /root/ -git clone https://github.com/THUDM/slime.git -cd slime/ -pip install -e . --no-deps -``` - -Download the model and data: - -```bash -# hf checkpoint -hf download zai-org/GLM-Z1-9B-0414 --local-dir /root/GLM-Z1-9B-0414 - -# train data -hf download --repo-type dataset zhuzilin/dapo-math-17k \ - --local-dir /root/dapo-math-17k - -# eval data -hf download --repo-type dataset zhuzilin/aime-2024 \ - --local-dir /root/aime-2024 -``` - -Convert the Hugging Face checkpoint to a Megatron-loadable Hugging Face checkpoint: - -```bash -# mcore checkpoint -cd /root/slime -source scripts/models/glm4-9B.sh -PYTHONPATH=/root/Megatron-LM python tools/convert_hf_to_torch_dist.py \ - ${MODEL_ARGS[@]} \ - --hf-checkpoint /root/GLM-Z1-9B-0414 \ - --save /root/GLM-Z1-9B-0414_torch_dist -``` - -## Run Training - -Execute the training: - -```bash -cd /root/slime -bash scripts/run-glm4-9B.sh -``` - -### Parameter Introduction - -Here, we will briefly introduce the various components of the [run-glm4-9B.sh](https://github.com/THUDM/slime/blob/main/scripts/run-glm4-9B.sh) script: - -#### MODEL\_ARGS - -```bash -SCRIPT_DIR="$(cd -- "$(dirname -- "${BASH_SOURCE[0]}")" &>/dev/null && pwd)" -source "${SCRIPT_DIR}/models/glm4-9B.sh" -``` - -Reads the model's config from [scripts/models/glm4-9B.sh](https://github.com/THUDM/slime/blob/main/scripts/models/glm4-9B.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/THUDM/slime/tree/main/scripts/models/). - -⚠️ Ensure that settings such as `--rotary-base` in the model configuration file match the settings of the model you are currently training. This is because different models, even with the same architecture, might use different values. If needed, you can override these parameters in your script after loading the model weights. For instance: - -```bash -source "${SCRIPT_DIR}/models/glm4-9B.sh" - -MODEL_ARGS += ( --rotary-base 10000 ) -``` - -#### CKPT\_ARGS - -```bash -CKPT_ARGS=( - # HF checkpoint required by vLLM; we also read the tokenizer from here - --hf-checkpoint /root/GLM-Z1-9B-0414 - # Checkpoint for the reference model - --ref-load /root/GLM-Z1-9B-0414_torch_dist - # Load directory for the actor; if empty, it will be loaded from `ref_load` - --load /root/GLM-Z1-9B-0414_slime/ - --save /root/GLM-Z1-9B-0414_slime/ - --save-interval 20 -) -``` - -#### ROLLOUT\_ARGS - -```bash -ROLLOUT_ARGS=( - # Prompt dataset, each line is a JSON object - --prompt-data /root/dapo-math-17k/dapo-math-17k.jsonl - --input-key prompt - --label-key label - # If the `input_key` in the prompt contains an OpenAI message, - # tokenizer.apply_chat_template(...) will be executed - --apply-chat-template - # Whether to shuffle the data - --rollout-shuffle - - # Reward model type. - # slime provides many types and --custom-rm-path for custom models - --rm-type deepscaler - - # Total number of rollouts to train - --num-rollout 3000 - # Number of prompts in one rollout - --rollout-batch-size 32 - # Number of responses to sample per prompt - # A rollout will have rollout_batch_size * n_samples_per_prompt items - --n-samples-per-prompt 8 - # Rollout sampling parameters - --rollout-max-response-len 8192 - --rollout-temperature 1 - - # Number of training steps corresponding to one rollout - --num-steps-per-rollout 1 - # Whether to balance data during training, which might improve speed - --balance-data -) -``` - -#### EVAL\_ARGS - -During evaluation, most rollout parameters are inherited, but we provide some parameters that can override the rollout configuration, allowing for different sampling strategies for training and evaluation. - -```bash -EVAL_ARGS=( - --eval-interval 5 - --eval-prompt-data /root/aime-2024/aime-2024.jsonl - --n-samples-per-eval-prompt 16 - --eval-max-response-len 16384 - --eval-top-p 1 -) -``` - -#### PERF\_ARGS - -A set of Megatron's parallelism parameters. Only `--use-dynamic-batch-size` and `--max-tokens-per-gpu` are added by slime. - -`max_tokens_per_gpu` specifies the maximum number of tokens each GPU can process. When `use_dynamic_batch_size` is enabled, it will try to pack data of varying lengths within a batch up to `max_tokens_per_gpu`, thus forming a dynamic micro-batch size. If a single data item's length exceeds `max_tokens_per_gpu`, it will form its own batch without being truncated. When context parallelism (CP) is enabled, it allows the CP GPUs to share data with a total length of `CP * max_tokens_per_gpu` tokens. - -When `dynamic_batch_size` is enabled, the traditional `micro_batch_size` is ignored. - -⚠️ slime 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 2 - --sequence-parallel - --pipeline-model-parallel-size 1 - --context-parallel-size 2 - --expert-model-parallel-size 1 - --expert-tensor-parallel-size 1 - - --recompute-granularity full - --recompute-method uniform - --recompute-num-layers 1 - - # --micro-batch-size 1 - --use-dynamic-batch-size - --max-tokens-per-gpu 4608 -) -``` - -#### GRPO\_ARGS - -Here are some GRPO-related parameters: - -```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 -) -``` - -#### OPTIMIZER\_ARGS - -```bash -OPTIMIZER_ARGS=( - --optimizer adam - --lr 1e-6 - --lr-decay-style constant - --weight-decay 0.1 - --adam-beta1 0.9 - --adam-beta2 0.98 -) -``` - -#### VLLM\_ARGS - -Parameters required by vLLM. Here, `--rollout-num-gpus-per-engine` corresponds to vLLM's `tp_size`. Other vLLM parameters are passed to slime by adding the `--vllm-` prefix. - -```bash -VLLM_ARGS=( - --rollout-num-gpus-per-engine 2 -) -``` - -⚠️ slime uses `vllm-router` to schedule multiple vLLM engines. - -### Co-located Training and Inference - -In the original script, the resource configuration is as follows: - -```bash -ray job submit ... \ - -- python3 train.py \ - --actor-num-nodes 1 \ - --actor-num-gpus-per-node 4 \ - --rollout-num-gpus 4 \ - ... -``` - -This enables decoupled training and inference, where the training part will use 1 machine with 4 GPUs, and the inference will use another 4 GPUs. - -If you want to use the co-located feature, you need to add `--colocate` and remove `--rollout-num-gpus`: - -```bash -ray job submit ... \ - -- python3 train.py \ - --actor-num-nodes 1 \ - --actor-num-gpus-per-node 8 \ - --colocate \ - ... -``` - -In this case, both training and inference will share these 8 GPUs. - -⚠️ When using co-located training and inference, Megatron will always occupy some GPU memory. Therefore, you need to adjust `--vllm-gpu-memory-utilization` to reduce the proportion of memory occupied by vLLM. - -### Dynamic Sampling - -slime supports more complex sampling schemes, such as the dynamic sampling in [DAPO](https://dapo-sia.github.io/). To enable dynamic sampling, you need to configure: - -```bash - --over-sampling-batch-size ${OVER_SAMPLING_BS} \ - --dynamic-sampling-filter-path \ - slime.rollout.filter_hub.dynamic_sampling_filters.check_reward_nonzero_std \ -``` - -Here, `over_sampling_batch_size` needs to be greater than `rollout_batch_size`. For example: - -```bash - --rollout-batch-size 32 \ - --n-samples-per-prompt 8 \ - --over-sampling-batch-size 64 \ -``` - -The sampling will then directly sample 64 prompts, with 8 samples per prompt. Since slime performs asynchronous sampling internally, we will receive the 8 responses for each prompt sequentially. Upon receiving responses, they will be filtered using the function specified by `dynamic_sampling_filter_path`. If they pass, these 8 data points are kept; otherwise, they are discarded. The function in the example checks if the answers are all correct or all incorrect: - -```python -def check_reward_nonzero_std(args, samples: list[Sample], **kwargs): - rewards = [sample.reward for sample in samples] - return torch.tensor(rewards, dtype=torch.float).std() > 0.0 -``` - -When we have received 32 \* 8 data points, we will immediately stop sampling and will not wait for the remaining data to be sampled. If more than 32 prompts' worth of data is discarded (leaving fewer than 32 prompts' worth), we will then sample another 64 prompts. - -### Partial Rollout - -During the process of dynamic sampling, a large number of requests are aborted prematurely. We can configure the `--partial-rollout` parameter to save these partially generated requests to a data buffer. In the next rollout, these requests can be retrieved to continue data generation, thereby further optimizing performance. - -You can customize how data is retrieved from the buffer by configuring the `--buffer-filter-path`. The default function is: - -```python -def pop_first(args, rollout_id, buffer: list[list[Sample]], num_samples: int) -> list[list[Sample]]: - num_to_pop = min(len(buffer), num_samples) - samples = buffer[:num_to_pop] - del buffer[:num_to_pop] - return samples -``` - -This means that each time, the data corresponding to the first `num_samples` prompts is retrieved, totaling `num_samples * n_samples_per_prompt` items. - -⚠️ The `sample.metadata` of each partial rollout sample stores the rollout ID from its initial generation, which can be used for data filtering. diff --git a/docs/en/examples/glm4.7-30B-A3B.md b/docs/en/examples/glm4.7-30B-A3B.md deleted file mode 100644 index 8e10043b5..000000000 --- a/docs/en/examples/glm4.7-30B-A3B.md +++ /dev/null @@ -1,151 +0,0 @@ -# GLM-4.7-Flash with 8×H100 - - -## Environment Preparation - -The environment setup, data, and checkpoint conversion are the same as for the Qwen3-4B model. You can refer to [Example: Qwen3-4B Model](qwen3-4B.md), replacing mentions of Qwen3-4B with GLM-4.7-Flash. - -### Download Model - -```bash -hf download THUDM/GLM-4.7-Flash --local-dir /root/GLM-4.7-Flash -``` - -### Convert Checkpoint - -To convert the Hugging Face checkpoint to torch_dist format: - -```bash -cd /root/slime -pip install -e . --no-deps -source scripts/models/glm4.7-30B-A3B.sh -PYTHONPATH=/root/Megatron-LM/ torchrun --nproc-per-node 8 \ - tools/convert_hf_to_torch_dist.py \ - ${MODEL_ARGS[@]} \ - --hf-checkpoint /root/GLM-4.7-Flash/ \ - --save /root/GLM-4.7-Flash_torch_dist/ -``` - -## Run Training - -Execute the training script: - -```bash -cd /root/slime -bash scripts/run-glm4.7-30B-A3B-8gpus.sh -``` - -### Parameter Introduction - -Here, we will briefly introduce the key parts in the [run-glm4.7-30B-A3B-8gpus.sh](https://github.com/THUDM/slime/blob/main/scripts/run-glm4.7-30B-A3B-8gpus.sh) script. - -#### MoE Configuration - -GLM-4.7-Flash is a Mixture-of-Experts (MoE) model with 64 routed experts (top-4 activation) and 1 shared expert. It has 47 layers: 1 dense layer + 46 MoE layers. - -1. To support running GLM-4.7-Flash on 8×H100, we need to enable Megatron's CPU Adam to save GPU memory: - - ```bash - OPTIMIZER_ARGS=( - ... - --optimizer-cpu-offload - --overlap-cpu-optimizer-d2h-h2d - --use-precision-aware-optimizer - ) - ``` - -2. Enable MoE optimization in Megatron. For single-node 8×H100, we use TP=1, EP=8: - - ```bash - PERF_ARGS=( - --tensor-model-parallel-size 1 - --pipeline-model-parallel-size 1 - --context-parallel-size 1 - --expert-model-parallel-size 8 - --expert-tensor-parallel-size 1 - ... - ) - ``` - -3. Enable MoE expert parallelism in vLLM. GLM-4.7-Flash is a non-MLA model, - so we use 8-way data parallelism on the attention block with expert - parallelism on the experts: - - ```bash - VLLM_ARGS=( - --rollout-num-gpus-per-engine 8 - --vllm-gpu-memory-utilization 0.7 - --vllm-data-parallel-size 8 - --vllm-enable-expert-parallel - ... - ) - ``` - -#### MTP Speculative Decoding (Inference Acceleration) - -GLM-4.7-Flash includes 1 MTP (Multi-Token Prediction) layer, which can be used for speculative decoding during inference to speed up rollout generation. To enable this, add the following to `VLLM_ARGS`: - -```bash -VLLM_ARGS=( - ... - # MTP speculative decoding (EAGLE) - --vllm-speculative-config '{"method":"mtp","num_speculative_tokens":3}' -) -``` - -This enables vLLM to use the model's MTP layer as a draft model for EAGLE-style speculative decoding. The MTP layer predicts multiple future tokens, and vLLM verifies them in parallel, leading to faster generation. - -> ⚠️ **Note**: Speculative decoding requires additional GPU memory. If you encounter OOM issues, try reducing `--vllm-gpu-memory-utilization` or disabling speculative decoding. - -#### MTP Training - -slime also supports training MTP layers jointly with the main model for models that have MTP weight conversion implemented (e.g., MiMo, GLM-4.7). When enabled, the relevant arguments are: - -```bash -# Add MTP layer count to model config -MODEL_ARGS+=(--mtp-num-layers 1) - -# Enable MTP training -SPEC_ARGS=( - --enable-mtp-training - --mtp-loss-scaling-factor 0.2 -) -``` - -- `--mtp-num-layers 1`: Tells Megatron to load the MTP layer from the checkpoint. -- `--enable-mtp-training`: Enables gradient computation for MTP layers. Without this flag, the MTP layer is loaded but frozen. -- `--mtp-loss-scaling-factor 0.2`: Weight of the MTP loss relative to the main policy loss. Default is 0.2. - -> **Note**: MTP training requires the MTP checkpoint bridge to properly convert weights between HuggingFace and Megatron formats. The `GLM4MoELiteBridge` (in `slime_plugins/mbridge/glm4moe_lite.py`) extends the DeepSeek V3 bridge with dynamic MTP layer indexing to support GLM-4.7-Flash's 47-layer architecture. -> -> For other models with MTP training support (e.g., MiMo), see `scripts/run-mimo-7B-rl-eagle.sh` as a reference. - -### Multi-Node Support - -For multi-node training (e.g., 2×8 H100), use the multi-node script: - -```bash -cd /root/slime -export BASE_DIR=/shared/path # accessible by all nodes -bash scripts/run-glm4.7-30B-A3B.sh -``` - -Key modifications for multi-node: - - - Place the model and data on a path accessible by all nodes. - - Set `MASTER_ADDR` to an address accessible by all nodes. - - Remove CPU Adam configurations (distributed optimizer reduces per-GPU memory usage). - - Adjust parallelism: e.g., TP=4, PP=2, EP=8, CP=2. - -When the total number of GPUs is not a multiple or divisor of the total number of experts (64), enable vLLM's EPLB (Expert Parallelism Load Balancer) and configure redundant experts via `--vllm-eplb-config`. For example, in a 24-GPU scenario: - -```bash -VLLM_ARGS=( - --rollout-num-gpus-per-engine 24 - --vllm-gpu-memory-utilization 0.7 - --vllm-data-parallel-size 3 - --vllm-enable-expert-parallel - --vllm-enable-eplb - --vllm-eplb-config '{"num_redundant_experts": 16}' -) -``` diff --git a/docs/en/examples/glm4.7-355B-A32B.md b/docs/en/examples/glm4.7-355B-A32B.md deleted file mode 100644 index 3ae3dc497..000000000 --- a/docs/en/examples/glm4.7-355B-A32B.md +++ /dev/null @@ -1,183 +0,0 @@ -# GLM-4.7 with 64xH100 - -## Environment Preparation - -The environment setup and dataset download are the same as for the Qwen3-4B model. You can refer to [Example: Qwen3-4B Model](qwen3-4B.md), replacing mentions of Qwen3-4B with GLM-4.7. - -### Prerequisites - -GLM-4.7 follows the standard slime Docker environment. For multi-node launches, make sure all nodes can access the same `$BASE_DIR` path and unset proxy variables before starting Ray workers: - -```bash -unset http_proxy https_proxy HTTP_PROXY HTTPS_PROXY -``` - -### Download Model - -```bash -hf download zai-org/GLM-4.7 --local-dir $BASE_DIR/GLM-4.7-355B-A32B -``` - -### Convert Checkpoint - -To convert the Hugging Face checkpoint to torch_dist format, use 2 nodes x 8 GPUs: - -```bash -cd /root/slime -pip install -e . --no-deps -source scripts/models/glm4.5-355B-A32B.sh -PYTHONPATH=/root/Megatron-LM/ torchrun \ - --nproc-per-node 8 \ - --master-addr ${MASTER_ADDR} --master-port 12345 \ - --nnodes=2 --node-rank ${NODE_RANK} \ - tools/convert_hf_to_torch_dist.py \ - ${MODEL_ARGS[@]} \ - --hf-checkpoint $BASE_DIR/GLM-4.7-355B-A32B/ \ - --save $BASE_DIR/GLM-4.7-355B-A32B_torch_dist/ -``` - -Here, `MASTER_ADDR` is the IP of node0, and `NODE_RANK` is the node index, configured just like a multi-node `torchrun` job. - -## Run Training - -Execute the training script from node0: - -```bash -cd /root/slime -export BASE_DIR=/shared/path # accessible by all nodes -bash scripts/run-glm4.7-355B-A32B.sh -``` - -### Parameter Introduction - -Here, we briefly introduce the key parts in the [run-glm4.7-355B-A32B.sh](https://github.com/THUDM/slime/blob/main/scripts/run-glm4.7-355B-A32B.sh) script. - -#### MoE Configuration - -GLM-4.7 is a Mixture-of-Experts (MoE) model with 160 routed experts (top-8 activation) and shared experts. It has 92 layers: 3 dense layers + 89 MoE layers. - -1. To support GLM-4.7 on 64xH100, we enable Megatron's CPU Adam to save GPU memory: - - ```bash - OPTIMIZER_ARGS=( - ... - --optimizer-cpu-offload - --overlap-cpu-optimizer-d2h-h2d - --use-precision-aware-optimizer - ) - ``` - -2. Enable MoE optimization in Megatron. For the provided 64xH100 example, we use TP=8, PP=4, CP=2, and EP=16: - - ```bash - PERF_ARGS=( - --tensor-model-parallel-size 8 - --sequence-parallel - --pipeline-model-parallel-size 4 - --context-parallel-size 2 - --expert-model-parallel-size 16 - --expert-tensor-parallel-size 1 - ... - --use-dynamic-batch-size - --max-tokens-per-gpu 16384 - ) - ``` - -3. Enable MoE expert parallelism in vLLM. GLM-4.7 is non-MLA, so we use 4-way - data parallelism on the attention block with expert parallelism on the - experts (EP size is auto-derived as `tensor_parallel_size × data_parallel_size`): - - ```bash - VLLM_ARGS=( - --rollout-num-gpus-per-engine 32 - --vllm-gpu-memory-utilization 0.7 - --vllm-data-parallel-size 4 - --vllm-enable-expert-parallel - ... - ) - ``` - -#### MTP Speculative Decoding (Inference Acceleration) - -GLM-4.7 includes MTP (Multi-Token Prediction) layers that can be used for speculative decoding during inference to speed up rollout generation. To enable this, add the following to `VLLM_ARGS`: - -```bash -VLLM_ARGS=( - ... - # MTP speculative decoding (EAGLE) - --vllm-speculative-config '{"method":"mtp","num_speculative_tokens":3}' -) -``` - -This lets vLLM use the model's MTP layer as the draft model for EAGLE-style speculative decoding. - -> ⚠️ **Note**: Speculative decoding requires additional GPU memory. If you encounter OOM issues, try reducing `--vllm-gpu-memory-utilization` or disabling speculative decoding. - -#### MTP Training - -slime also supports training the MTP layers jointly with the main model for GLM-4.7. When enabled, the relevant arguments are: - -```bash -# Add MTP layer count to model config -MODEL_ARGS+=(--mtp-num-layers 1) - -# Enable MTP training -MTP_ARGS=( - --enable-mtp-training - --mtp-loss-scaling-factor 0.2 -) -``` - -- `--mtp-num-layers 1`: Tells Megatron to load the MTP layer from the checkpoint. -- `--enable-mtp-training`: Enables gradient computation for MTP layers. Without this flag, the MTP layer is loaded but frozen. -- `--mtp-loss-scaling-factor 0.2`: Weight of the MTP loss relative to the main policy loss. Default is 0.2. - -> **Note**: MTP training for GLM-4.7 relies on `GLM4MoEBridge` (in `slime_plugins/mbridge/glm4moe.py`) to map regular and MTP weights between HuggingFace and Megatron formats. - -#### Multi-Node Support - -This example already targets multi-node training. Before launching: - -- Place the model checkpoints and datasets on a path accessible by all nodes. -- Set `MASTER_ADDR` to an address reachable by all nodes. -- Unset proxy variables before starting Ray workers. -- Provide a `HOSTFILE` listing worker IPs (one per line) and export `HOSTFILE=/path/to/hostfile` before launching. -- Adjust parallelism coherently. The default example uses TP=8, PP=4, EP=16, CP=2, while rollout uses 32 GPUs per engine with DP=4 + expert parallel. - -If your rollout GPU count does not divide the expert count cleanly, enable vLLM's EPLB (Expert Parallelism Load Balancer) and configure redundant experts via `--vllm-eplb-config`, e.g.: - -```bash ---vllm-enable-eplb ---vllm-eplb-config '{"num_redundant_experts": 16}' -``` - -## FP8 Rollout - -The open-source FP8 checkpoint of GLM-4.7 uses per-channel quantization, which cannot currently enable DeepEP in vLLM. You can convert it to a 128x128 per-block FP8 checkpoint with the tool provided in slime: - -```bash -cd /root/slime -python tools/convert_hf_to_fp8.py \ - --model-dir $BASE_DIR/GLM-4.7-355B-A32B/ \ - --save-dir $BASE_DIR/GLM-4.7-355B-A32B-FP8/ \ - --strategy block --block-size 128 128 \ - --max-workers 4 -``` - -Then switch `--hf-checkpoint` to `$BASE_DIR/GLM-4.7-355B-A32B-FP8/` to enable FP8 rollout. - -An example FP8 `VLLM_ARGS` setup is: - -```bash -VLLM_ARGS=( - --rollout-num-gpus-per-engine 32 - --vllm-gpu-memory-utilization 0.7 - --vllm-data-parallel-size 32 - --vllm-enable-expert-parallel - --vllm-cudagraph-capture-sizes 1 2 4 8 $(seq 16 8 128) - - --vllm-speculative-config '{"method":"mtp","num_speculative_tokens":3}' - - --vllm-all2all-backend deepep_high_throughput -) -``` diff --git a/docs/en/examples/qwen3-4b-base-openhermes.md b/docs/en/examples/qwen3-4b-base-openhermes.md deleted file mode 100644 index 55ea305f8..000000000 --- a/docs/en/examples/qwen3-4b-base-openhermes.md +++ /dev/null @@ -1,86 +0,0 @@ -# SFT Qwen3-4B-Base - - -## Environment Preparation - -First, we need to create a mirror environment and convert the `Qwen3-4B-Base` model by following the [Example: Qwen3-4B Model](qwen3-4B.md). - -After that, we will process the SFT data. Here, we use the classic [OpenHermes-2.5](https://huggingface.co/datasets/teknium/OpenHermes-2.5) as an example. First, we process the data into a format suitable for `slime` to load. You can use the following script to add a column that conforms to the OpenAI message format and save it to `/root/openhermes2_5.parquet`. - -```python -from datasets import load_dataset - -ds = load_dataset("teknium/OpenHermes-2.5")["train"] - -def convert(sample): - conversations = sample["conversations"] - - def convert_role(role): - if role == "human": - return "user" - elif role == "gpt": - return "assistant" - elif role == "system": - return "system" - else: - raise ValueError(f"Unknown role: {role}") - - messages = [ - { - "role": convert_role(turn["from"]), - "content": turn["value"], - } - for turn in conversations - ] - - return {"messages": messages} - -ds = ds.map(convert) -ds.to_parquet("/root/openhermes2_5.parquet") -``` - -## Execute Training - -Execute the training: - -```bash -cd /root/slime -bash script/run-qwen3-4B-base-sft.sh -``` - -### Parameter Introduction - -You can compare [run-qwen3-4B-base-sft.sh](https://github.com/THUDM/slime/blob/main/scripts/run-qwen3-4B-base-sft.sh) with [run-qwen3-4B.sh](https://github.com/THUDM/slime/blob/main/scripts/run-qwen3-4B.sh). You will find that besides changing the model from the instruct version to the base model, the main adjustments are as follows: - -1. Removed `VLLM_ARGS` and `GRPO_ARGS`. This is because it is not necessary to start vLLM or configure GRPO-related settings during the SFT process. - -2. Renamed `ROLLOUT_ARGS` to `SFT_ARGS` and configured it as follows: - - ```bash - SFT_ARGS=( - --rollout-function-path slime.rollout.sft_rollout.generate_rollout - --prompt-data /root/openhermes2_5.parquet - --input-key messages - --rollout-shuffle - --num-epoch 3 - --rollout-batch-size 128 - --global-batch-size 128 - - --loss-type sft_loss - --calculate-per-token-loss - --disable-compute-advantages-and-returns - --debug-train-only - ) - ``` - - SFT actually reuses the custom rollout functionality of slime. By using `--rollout-function-path`, the data generation part is switched from the RL rollout that uses `vLLM` to the SFT version that reads data from a file, which is `slime.rollout.sft_rollout.generate_rollout`. - - For SFT, it is recommended to set `rollout_batch_size` and `global_batch_size` to the same value and not to configure `n_samples_per_prompt`. This is equivalent to training one batch right after reading one batch. - - `slime` also supports different loss types, and we configure the SFT loss using `--loss-type sft_loss`. - - As for `--calculate-per-token-loss`, this is because `slime` defaults to calculating the per-sample mean for GRPO. In general SFT training, the average is taken over all unmasked tokens in a batch, so it is recommended to configure this. - - Finally, `--disable-compute-advantages-and-returns` indicates that there is no need to pre-calculate log probabilities during the SFT process, and `--debug-train-only` means that `vLLM` does not need to be initialized. - -3. Used `train_async.py` instead of `train.py`. This is to leverage the asynchronous training process to implement data prefetching. diff --git a/docs/en/get_started/quick_start.md b/docs/en/get_started/quick_start.md index d993eaa69..98560198a 100644 --- a/docs/en/get_started/quick_start.md +++ b/docs/en/get_started/quick_start.md @@ -53,7 +53,7 @@ You can download required models and datasets from platforms like Hugging Face, ```bash # Download model weights (GLM-Z1-9B) -hf download zai-org/GLM-Z1-9B-0414 --local-dir /root/GLM-Z1-9B-0414 +hf download zai-org/Qwen3-4B --local-dir /root/Qwen3-4B # Download training dataset (dapo-math-17k) hf download --repo-type dataset zhuzilin/dapo-math-17k \ @@ -70,11 +70,11 @@ hf download --repo-type dataset zhuzilin/aime-2024 \ When using Megatron as the training backend, you need to first convert Hugging Face format model weights to Megatron `torch_dist` format. -First, load the configuration file of the target model. The `slime/scripts/models` directory contains configuration files for supported models. You need to `source` the corresponding model script to load the configuration parameters into the current environment. Here we use GLM4-9B model as an example, and it's similar for Qwen3-4B, Qwen3.5, Qwen3.6, GLM-4.7-Flash, Qwen3-30B-A3B, etc. +First, load the configuration file of the target model. The `slime/scripts/models` directory contains configuration files for supported models. You need to `source` the corresponding model script to load the configuration parameters into the current environment. Here we use Qwen3-4B model as an example. ```bash cd /root/slime -source scripts/models/glm4-9B.sh +source scripts/models/qwen3-4B.sh ``` Next, run the conversion script. Please note the following parameters: @@ -84,12 +84,11 @@ Next, run the conversion script. Please note the following parameters: ```bash PYTHONPATH=/root/Megatron-LM python tools/convert_hf_to_torch_dist.py \ ${MODEL_ARGS[@]} \ - --hf-checkpoint /root/GLM-Z1-9B-0414 \ - --save /root/GLM-Z1-9B-0414_torch_dist + --hf-checkpoint /root/Qwen3-4B \ + --save /root/Qwen3-4B_torch_dist ``` For larger models, you can use `torchrun` to start the conversion script to convert with multi-gpus or even multi-nodes. -Note: When converting the kimi-k2 model weights, you need to open config.json in the model path and change "model_type": "kimi_k2" to "model_type": "deepseek_v3". ### Convert from Megatron Format to Hugging Face Format @@ -98,8 +97,8 @@ You can use the following script to convert the saved Megatron checkpoints back ```bash PYTHONPATH=/root/Megatron-LM python tools/convert_torch_dist_to_hf.py \ --input-dir /path/to/torch_dist_ckpt/iter_xxx/ \ - --output-dir /root/GLM-Z1-9B-0414-iter_xxx \ - --origin-hf-dir /root/GLM-Z1-9B-0414 + --output-dir /root/Qwen3-4B-iter_xxx \ + --origin-hf-dir /root/Qwen3-4B ``` Note that as Megatron will do padding to embedding for better performance, it may happen that the converted embedding is not correct. In that case, please manually set `--vocab-size` during convertion. @@ -110,24 +109,24 @@ After completing the above preparation work, you can run the training script. ```bash cd /root/slime -bash scripts/run-glm4-9B.sh +bash scripts/run-qwen3-4B.sh ``` -We still use the run-glm4-9B.sh script as an example to briefly analyze the main parameters. +We still use the run-qwen3-4B.sh script as an example to briefly analyze the main parameters. ### MODEL_ARGS: Model Configuration Parameters ```bash SCRIPT_DIR="$(cd -- "$(dirname -- "${BASH_SOURCE[0]}")" &>/dev/null && pwd)" -source "${SCRIPT_DIR}/models/glm4-9B.sh" +source "${SCRIPT_DIR}/models/qwen3-4B.sh" ``` -This part loads model configuration from the `scripts/models/glm4-9B.sh` file through the `source` command. These configurations are all hyperparameters required by Megatron. Since Megatron cannot directly read model configuration from checkpoints, it needs to be manually specified. We provide configuration examples for some commonly used models in the `scripts/models/` directory. +This part loads model configuration from the `scripts/models/qwen3-4B.sh` file through the `source` command. These configurations are all hyperparameters required by Megatron. Since Megatron cannot directly read model configuration from checkpoints, it needs to be manually specified. We provide configuration examples for some commonly used models in the `scripts/models/` directory. > ⚠️ **Note**: > Please make sure to check whether the parameters in the model configuration file (such as `--rotary-base`) completely match the model you are currently using. Different versions of the same model structure may use different configuration values. If you need to modify, you can directly override after `source`, for example: > ```bash -> source "${SCRIPT_DIR}/models/glm4-9B.sh" +> source "${SCRIPT_DIR}/models/qwen3-4B.sh" > MODEL_ARGS+=(--rotary-base 10000) > ``` @@ -136,14 +135,14 @@ This part loads model configuration from the `scripts/models/glm4-9B.sh` file th ```bash CKPT_ARGS=( # To load tokenizer and other information, won't actually use model weight parameters from hf path - --hf-checkpoint /root/GLM-Z1-9B-0414 + --hf-checkpoint /root/Qwen3-4B # Reference Model's Megatron format checkpoint - --ref-load /root/GLM-Z1-9B-0414_torch_dist + --ref-load /root/Qwen3-4B_torch_dist # Actor model loading path. Should typically match --save for checkpoint resumption # If empty or doesn't contain a valid checkpoint, loads from --ref-load instead - --load /root/GLM-Z1-9B-0414_slime/ + --load /root/Qwen3-4B_slime/ # Model save path during training - --save /root/GLM-Z1-9B-0414_slime/ + --save /root/Qwen3-4B_slime/ # Model save interval (steps) --save-interval 20 ) @@ -577,9 +576,6 @@ 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}') ``` -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: +slime has been deeply optimized for distributed training of large-scale Mixture of Experts (MoE) models. We provide an end-to-end training case for reference: -- [Example: 8xH100 Training GLM-4.7-Flash](../examples/glm4.7-30B-A3B.md) -- [Example: 64xH100 Training GLM-4.7](../examples/glm4.7-355B-A32B.md) -- [Example: 128xH100 Training DeepSeek-R1](../examples/deepseek-r1.md) -- The scripts such as `scripts/run_qwen3_30b_a3b.py`, `scripts/run_glm45_355b_a32b.py` also support multi-node training, though there are little documentations about it currently. +- [Example: Qwen3-30B-A3B with 8xH100](../examples/qwen3-30B-A3B.md) diff --git a/docs/en/index.rst b/docs/en/index.rst index c74814e41..b63632636 100644 --- a/docs/en/index.rst +++ b/docs/en/index.rst @@ -26,24 +26,18 @@ slime is the RL-framework behind GLM-4.7, GLM-4.6 and GLM-4.5. Apart from models :caption: Dense examples/qwen3-4B.md - examples/glm4-9B.md .. toctree:: :maxdepth: 1 :caption: MoE - examples/glm4.7-30B-A3B.md examples/qwen3-30B-A3B.md - examples/glm4.7-355B-A32B.md - examples/deepseek-r1.md .. toctree:: :maxdepth: 1 :caption: Advanced Features - advanced/on-policy-distillation.md advanced/speculative-decoding.md - advanced/low-precision.md advanced/reproducibility.md advanced/fault-tolerance.md advanced/pd-disaggregation.md @@ -55,7 +49,6 @@ slime is the RL-framework behind GLM-4.7, GLM-4.6 and GLM-4.5. Apart from models :maxdepth: 1 :caption: Other Usage - examples/qwen3-4b-base-openhermes.md _examples_synced/search-r1/README.md _examples_synced/fully_async/README.md _examples_synced/retool/README.md @@ -69,4 +62,3 @@ slime is the RL-framework behind GLM-4.7, GLM-4.6 and GLM-4.5. Apart from models developer_guide/debug.md developer_guide/trace.md developer_guide/profiling.md - diff --git a/docs/zh/advanced/low-precision.md b/docs/zh/advanced/low-precision.md deleted file mode 100644 index 27affe47f..000000000 --- a/docs/zh/advanced/low-precision.md +++ /dev/null @@ -1,120 +0,0 @@ -# 低精度训练 - -- [FP8 推理与 BF16 训练](#FP8-推理与-BF16-训练) -- [FP8 推理与 FP8 训练](#FP8-推理与-FP8-训练) -- [INT4 QAT 训练](#INT4-QAT-训练) - -## FP8 推理与 BF16 训练 - -你可以通过在 `--hf-checkpoint` 中设置块缩放（blockwise）量化的 HuggingFace 权重来运行 FP8 推演。转换命令如下： - -```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 -``` - -请确保转换后的权重目录中的 `config.json` 包含正确的 `quantization_config`，以便 slime 在权重更新期间自动使用 FP8 量化。 - -## FP8 推理与 FP8 训练 - -我们观察到，在训练和推理阶段同时使用 FP8，可以获得更高效的推理吞吐量，并降低训推不一致，从而使训练更稳定。更多细节请参考 [此博客](https://lmsys.org/blog/2025-11-25-fp8-rl/)。 - -### 快速开始 - -1. 使用上述 `tools/convert_hf_to_fp8.py` 将 HuggingFace 模型权重转换为 FP8 格式。 -2. 对于训练任务，需要添加以下参数： -```bash ---fp8-format e4m3 ---fp8-recipe blockwise -# --fp8-param-gather # [可选] 目前与 CPU Adam 优化器不兼容 - -``` - -同时，确保启用了环境变量 `NVTE_FP8_BLOCK_SCALING_FP32_SCALES`，目前我们会默认将这个参数设置为 `1`。 - -注意：目前只有 TransformerEngine 中的 `Linear` 和 `GroupLinear` 层使用 FP8 格式。`embedding` 和 `lm_head` 仍保持原始精度。如果未开启 `--fp8-param-gather`，TransformerEngine 中的权重将以 BF16 格式存储，仅在 `GEMM` 或 `GroupGEMM` 运算期间临时转换为 FP8。 - -3. 启动训练： - -```bash -# Qwen3-4B Int4 training -bash scripts/low_precision/run-qwen3-4b-fp8.sh - -# Qwen3-30B-A3B (2 nodes) -bash scripts/low_precision/run-qwen3-30b-a3b-fp8.sh -``` - -4. 使用保存的 ckpt：TransformerEngine 不会专门保存 FP8 量化后的权重；保存的 `torch_dist` ckpt 仍为原始精度（通常是 BF16）。如果你想在 FP8 下进行评估，需要先将 `torch_dist` 转换为 HuggingFace 格式，然后再转换为 FP8 HuggingFace 格式。 - -### 原理简述 - -以下是 slime 中 FP8 训练目前的实现方式： - -1. **初始化**：如果启用了 FP8 方案，相关层将在 FP8 上下文中构建。 -2. **训练过程**：在训练期间，权重和激活值会在线量化为 `nvfp8` 格式，并在前向和反向传播中调用 `cuBLAS FP8 GEMM` 进行计算。 -3. **权重更新**：在强化学习（RL）权重更新期间，Megatron 首先将 FP8 权重反量化为 BF16 格式，然后 slime 再将这些 BF16 权重重新量化为 FP8 并发送给 vLLM。（这种"反量化+再量化"的操作虽然不够优雅，但为了框架兼容性，目前尚未修改接口。） -4. **保存 ckpt**：与权重更新类似，从训练引擎保存 ckpt 时，也会反量化回 BF16 并以 `torch_dist` 格式保存。 - -### 待办事项 (TODO) - -目前 FP8 功能尚不完全成熟，仍存在以下已知问题： - -* FP8 权重存储（`--fp8-param-gather`）虽然能节省显存，但目前必须配合 TransformerEngine 的 `FusedAdam` 使用，这与 Megatron-LM 中的 CPU Adam 技术冲突。 - -## INT4 QAT 训练 - -本指南提供了 INT4 STE（直通估计器，Straight-Through Estimator）训练和 INT4 推理的示例。使用 INT4 推理可显著提升吞吐量，从而加速整个训练流水线（特别是在 rollout 生成阶段）。 - -### 快速开始 - -1. **将 HuggingFace 权重转换为 INT4** -首先，从 HuggingFace 下载 PTQ（训练后量化）校准数据集： -[wikitext-2-raw-v1](https://huggingface.co/datasets/Salesforce/wikitext/tree/main/wikitext-2-raw-v1) -接着，使用 `tools/convert_hf_to_int4.py` 脚本进行转换。确保 `--hf-checkpoint` 指向的目录中 `config.json` 包含正确的 `quantization_config`。 -```bash -python tools/convert_hf_to_int4.py \ - --input-dir /path/to/your/original/models \ - --output-dir /path/to/your/save/models \ - --data-dir /path/to/your/wikitext - -``` - -**提示**：如果你只想运行 INT4 推演（Rollout），只需将 `--hf-checkpoint` 设置为转换后的 INT4 路径即可。 -2. **启动 INT4 QAT 训练** -你需要配置特定的环境变量来设定量化参数。 -**环境变量说明：** -* **`OPEN_TRAINING_INT4_FAKE_QAT_FLAG`**: 启用 INT4 训练的伪量化（Fake Quantization）操作。 -* **`OPEN_TRAINING_INT4_GROUP_SIZE`**: 指定模型量化的块大小（Group Size）。 -* `moonlight-16B-A3B`、`qwen3-30B-A3B` 和 `qwen3-235B-A22B-int4` 设置为 **128**。 -* `kimi-k2-Thinking-int4` 设置为 **32**。 - -**配置示例：** -```json -RUNTIME_ENV_JSON="{ - \"env_vars\": { - ... - \"OPEN_TRAINING_INT4_FAKE_QAT_FLAG\": \"1\", - \"OPEN_TRAINING_INT4_GROUP_SIZE\": \"128\" - } -}" -``` - -**启动命令：** -```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 -``` - -*对于多节点环境，请根据您的集群配置启动 Ray 服务。* diff --git a/docs/zh/advanced/on-policy-distillation.md b/docs/zh/advanced/on-policy-distillation.md deleted file mode 100644 index 883601cc3..000000000 --- a/docs/zh/advanced/on-policy-distillation.md +++ /dev/null @@ -1,111 +0,0 @@ -# 在策略蒸馏 (On-Policy Distillation) - -在策略蒸馏 (OPD) 让学生模型在自己的 rollout 数据上训练，同时匹配教师模型的 token 级 log-probability，从而实现从大模型到小模型的知识传递。OPD 与 advantage estimator 正交——它作为 KL 惩罚项叠加在任意 estimator（GRPO、PPO、REINFORCE++ 等）之上。 - -## 关键参数 - -| 参数 | 说明 | -|------|------| -| `--use-opd` | 启用在策略蒸馏。使用 OPD 的必需标志。 | -| `--opd-type` | OPD 类型：`vllm` 或 `megatron`。启用 `--use-opd` 时必须设置。 | -| `--opd-kl-coef` | OPD KL 惩罚系数（默认值：1.0）。控制蒸馏信号相对于 RL advantage 的权重。 | -| `--opd-teacher-load` | 教师模型的 Megatron checkpoint 路径。`--opd-type=megatron` 时**必须**设置，`--opd-type=vllm` 时**不可**设置。 | -| `--opd-teacher-ckpt-step` | 可选的教师模型 checkpoint 步数。 | - -## 原理 - -OPD 通过减去一个 KL 惩罚项来修改 advantage 计算，鼓励学生匹配教师的输出分布： - -$$ -\hat{A}_t = A_t - \lambda_{\text{opd}} \cdot D_{\text{KL}}(P_{\text{teacher}} \| P_{\text{student}})_t -$$ - -其中 $A_t$ 是基础 estimator（如 GRPO）的原始 advantage，$\lambda_{\text{opd}}$ 是 `--opd-kl-coef`，$D_{\text{KL}}$ 是 token 级的逆 KL 散度。 - -因此 OPD 可以与任何 advantage estimator 组合使用，包括 GRPO、PPO、REINFORCE++ 和 GSPO。 - -## 两种教师模式 - -### vLLM 模式 (`--opd-type vllm`) - -教师模型运行在外部 vLLM 服务器上，教师的 log-probs 在 rollout 阶段获取。 - -**适用场景**：教师与学生架构不同，或教师模型太大无法与训练模型同时加载。 - -**工作流程**： -1. 外部 vLLM 服务器运行教师模型。 -2. 在 rollout 阶段，自定义 reward 函数（`slime.rollout.on_policy_distillation.reward_func`）将每个样本发送到教师服务器以获取 token 级 log-probs。 -3. 自定义后处理函数（`slime.rollout.on_policy_distillation.post_process_rewards`）将教师 log-probs 裁剪到 response 范围并存储到 `sample.teacher_log_probs` 中。 -4. 在训练阶段，从存储的教师 log-probs 计算 KL 惩罚并应用到 advantages 上。 - -**配置**： -```bash ---use-opd ---opd-type vllm ---opd-kl-coef 1.0 ---custom-rm-path slime.rollout.on_policy_distillation.reward_func ---custom-reward-post-process-path slime.rollout.on_policy_distillation.post_process_rewards ---rm-url http://:/inference/v1/generate -``` - -### Megatron 模式 (`--opd-type megatron`) - -教师模型通过 `--opd-teacher-load` 直接加载到 Megatron 中，教师的 log-probs 在训练前向传播阶段计算。 - -**适用场景**：教师与学生/参考模型架构相同，且能放入 GPU 显存。 - -**工作流程**： -1. 教师模型在初始化时作为额外的 Megatron 模型加载。 -2. 在训练前向传播阶段，教师模型为每个样本计算 log-probs。 -3. 内联计算 KL 惩罚并应用到 advantages。 - -**配置**： -```bash ---use-opd ---opd-type megatron ---opd-kl-coef 1.0 ---opd-teacher-load /path/to/teacher_torch_dist -``` - -> **注意**：教师 checkpoint 必须是 Megatron 格式（`torch_dist` 或 `torch`）。可以使用 `tools/convert_hf_to_torch_dist.py` 从 HuggingFace 格式转换。 - -## 运行示例 - -完整的示例脚本在 `examples/on_policy_distillation/` 中： - -### vLLM 教师 - -```bash -# 1. 下载模型和数据 -hf download Qwen/Qwen3-32B --local-dir /root/Qwen3-32B -hf download Qwen/Qwen3-8B --local-dir /root/Qwen3-8B -hf download --repo-type dataset zhuzilin/dapo-math-17k --local-dir /root/dapo-math-17k - -# 2. 转换学生模型 -cd /root/slime -source scripts/models/qwen3-8B.sh -PYTHONPATH=/root/Megatron-LM python tools/convert_hf_to_torch_dist.py \ - ${MODEL_ARGS[@]} \ - --hf-checkpoint /root/Qwen3-8B \ - --save /root/Qwen3-8B_torch_dist - -# 3. 运行 -bash examples/on_policy_distillation/run-qwen3-8B-opd.sh -``` - -### Megatron 教师 - -```bash -# 1. 将学生和教师模型都转换为 Megatron 格式 -# 2. 运行 -bash examples/on_policy_distillation/run-qwen3-8B-opd-megatron.sh -``` - -## 初步结果 - -使用 Qwen3-8B-Base 模型在 [OpenThoughts3-1.2M](https://huggingface.co/datasets/open-thoughts/OpenThoughts3-1.2M) 数据集的一部分上进行 SFT，然后在剩余数据上用 Qwen3-32B 教师进行在策略蒸馏，Math500 评测结果如下： - -| | Pass@1 | -|-----------------------------------------------|--------| -| Qwen3-8B-Base + SFT | 76% | -| Qwen3-8B-Base + SFT + On-Policy Distillation | 94% | diff --git a/docs/zh/developer_guide/ci.md b/docs/zh/developer_guide/ci.md index f6378e01a..f1993b470 100644 --- a/docs/zh/developer_guide/ci.md +++ b/docs/zh/developer_guide/ci.md @@ -20,7 +20,7 @@ slime 使用 GitHub Actions 进行 CI。测试通过 **PR label** 触发—— | Label | Job | 说明 | |---|---|---| | `run-ci-short` | `e2e-test-short` | Qwen2.5-0.5B 轻量级冒烟测试（4 GPU），用于快速反馈。 | -| `run-ci-megatron` | `e2e-test-megatron` | 核心 Megatron 训练测试，覆盖 Dense、MoE、PPO、MTP、OPD 等。 | +| `run-ci-megatron` | `e2e-test-megatron` | 核心 Megatron 训练测试，覆盖 Dense、MoE、PPO、OPD 等。 | | `run-ci-precision` | `e2e-test-precision` | 数值精度校验（并行一致性检查）。 | | `run-ci-ckpt` | `e2e-test-ckpt` | Checkpoint 保存/加载正确性（同步和异步保存）。 | | `run-ci-image` | `e2e-test-image` | 在 `inferactinc/public:vime-vllm-cu129-latest` 镜像上运行**全部**测试（用于镜像验证）。 | @@ -57,7 +57,7 @@ slime 使用 GitHub Actions 进行 CI。测试通过 **PR label** 触发—— - Dense 模型：GLM4-9B、Qwen3-4B（PPO） - MoE 模型：Qwen3-30B-A3B（DeepEP + FP8）、Qwen3.6-35B-A3B PD + Mooncake、Moonlight-16B-A3B -- 特殊场景：MiMo-7B MTP、Qwen2.5-0.5B debug rollout-then-train、OPD（vLLM teacher 模式） +- 特殊场景：Qwen2.5-0.5B debug rollout-then-train、OPD（vLLM teacher 模式） 所有测试使用 8 张 GPU。如果你正在修改 Megatron 训练逻辑、loss 计算或 checkpoint 转换，应该使用这个 label。 diff --git a/docs/zh/examples/deepseek-r1.md b/docs/zh/examples/deepseek-r1.md deleted file mode 100644 index 237540040..000000000 --- a/docs/zh/examples/deepseek-r1.md +++ /dev/null @@ -1,202 +0,0 @@ -# 128xH100 训练 DeepSeek R1 - -这里是使用 128xH100 进行 DeepSeek R1 RL 训练的示例。 - -我们会使用 bf16 进行训练，128x128 blockwise quant 的 fp8 格式进行推理，模型最大回复长度为 32k，并训练中会使用 dynamic sampling 对数据进行筛选。 - -在并行上，vLLM 方面我们会启用 expert parallelism 与 data parallelism；megatron 部分我们采用 tp8、pp4、ep32、cp4。 - -⚠️ 为了节省 GPU 显存，我们会使用 CPU Adam，每个 node（8xH100）会占用 1.4~1.5B 内存。如果单机的内存不够，可以通过增加 GPU，扩大并行的方式解决。 - -## 环境准备 - -搭建环境与下载数据的方法可以参考 [示例：Qwen3-4B](qwen3-4B.md)。 - -准备 DeepSeek R1 的 ckpt 首先需要在多机均可访问到的地址（下记为 `$BASE_DIR`）上下载 DeepSeek-R1： - -```bash -hf download deepseek-ai/DeepSeek-R1 --local-dir $BASE_DIR/DeepSeek-R1 -``` - -DeepSeek-R1 的 huggingface ckpt 为 block-quant 的 fp8 格式，为了转换一个 Megatron 可以加载的 torch dist 格式，需要先转化一个 bf16 的 huggingface ckpt： - -```bash -cd slime/ -python tools/fp8_cast_bf16.py --input-fp8-hf-path $BASE_DIR/DeepSeek-R1 --output-bf16-hf-path $BASE_DIR/DeepSeek-R1-bf16/ -``` - -之后我们需要将 bf16 版本的 DeepSeek-R1 转换为 torch dist 格式。具体为在 4 台机器上分别执行： - -```bash -cd slime/ -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/ -``` - -其中 `MASTER_ADDR` 为 node0 的 ip，`NODE_RANK` 表示这是第几台机器，这两者就像是在多机 `torchrun` 的时候进行的配置。 - -## 执行训练 - -在 node0 运行： - -```bash -cd slime/ -bash scripts/run-deepseek-r1.sh -``` - -在其他 node 需要通过如下的指令加入 ray 集群： - -```bash -ray start --address=${MASTER_ADDR}:6379 --num-gpus 8 --node-ip-address ${WORKER_IP} --disable-usage-stats" -``` - -或者如果你能获取到所有节点的 ip 列表，例如有一个 mpi hostfie（每一行为 `ip slot=8`），那么可以在 `scripts/run-deepseek-r1.sh` 中的 `ray start --head` 指令之后加入如下的指令，从而只需要从 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 -``` - -### 参数简介 - -```bash -SCRIPT_DIR="$(cd -- "$(dirname -- "${BASH_SOURCE[0]}")" &>/dev/null && pwd)" -source "${SCRIPT_DIR}/models/deepseek-v3.sh" -``` - -从 [scripts/models/deepseek-v3.sh](https://github.com/THUDM/slime/blob/main/scripts/models/deepseek-v3.sh) 读取模型的 config。这些 config 都是 megatron 的参数。在使用 megatron 进行训练的时候，megatron 无法从 ckpt 中读取模型 config，需要我们自行配置。我们在 [scripts/models](https://github.com/THUDM/slime/tree/main/scripts/models/) 中提供了一些样例。 - - -#### CKPT_ARGS - -```bash -CKPT_ARGS=( - # vLLM 需要的 hf ckpt，我们也会从这里读 tokenizer - --hf-checkpoint $BASE_DIR/DeepSeek-R1/ - #--hf-checkpoint $BASE_DIR/DeepSeek-R1-bf16/ - --ref-load $BASE_DIR/DeepSeek-R1_torch_dist/ - # actor 的 load dir，如果是空的，会从 `ref_load` 里面读 - --load $BASE_DIR/DeepSeek-R1_slime/ - --save $BASE_DIR/DeepSeek-R1_slime/ - --save-interval 20 -) -``` - -slime 会根据 `hf_checkpoint` 中的量化配置从而在训练中进行在线量化。例如当前的例子中，我们使用的是 DeepSeek R1 的 fp8 ckpt，那么在进行参数更新的时候，我们会首先将参数进行 blockwise quant，再传至 vLLM。 - -#### PERF_ARGS - -一堆 megatron 的并行参数，只有 `--use-dynamic-batch-size` 与 `--max-tokens-per-gpu` 是 slime 添加的。 - -megatron 的部分，我们配置了 tp8、pp4、cp4、ep32，由于 DeepSeek-R1 有 61 层，不能被 4 整除，所以我们专门配置最后一个 pp stage 为 13 层。 - -`max_tokens_per_gpu` 是指每张卡最多跑多少 token，在开启 `use_dynamic_batch_size` 之后，会尽可能将一个 batch 内部长短不一的数据拼到 `max_tokens_per_gpu`，从而组成动态的 micro batch size，如果有一条数据长度超过了 `max_tokens_per_gpu`，则自成一条，不会对数据进行截断。在开启 context parallel (CP) 时，会让 CP 张卡去上的数据去共享总长为 `CP * max_tokens_per_gpu` 的 token。 - -在开启 dynamic_batch_size，会忽略传统的 `micro_batch_size`。 - -⚠️ slime 总是会通过 data packing 的方法训练模型，并且严格保证 per sample loss 或 per token loss，也就是开启 dynamic batch size 不会对 loss 计算有影响，推荐开启。 - -```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 - -目前 slime 这是一些 grpo 相关的参数： - -```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 -) -``` - -如果希望训练时不加载 reference model，需要去掉 `--use-kl-loss` 并设置 `--kl-coef 0.00`（默认值为 0）。 - -#### OPTIMIZER_ARGS - -我们通过了如下几个参数配置了 CPU Adam，用来节省显存。 - -```bash -OPTIMIZER_ARGS=( - ... - - --optimizer-cpu-offload - --overlap-cpu-optimizer-d2h-h2d - --use-precision-aware-optimizer -) -``` - -#### VLLM_ARGS - -vLLM 所需的参数，这里 `--rollout-num-gpus-per-engine` 对应 vLLM 的 `tp_size`，除此之外的 vLLM 参数均通过添加 `--vllm-` 的前缀来传给 slime。 - -`--vllm-server-concurrency` 是 slime 的特有参数，用于防止同时发给 vLLM 引擎的并发太大打爆 HTTP server，默认为 512。但是我们现在是 8 机一个 server，为了保证每个 dp rank 能有 128 的并发，我们调整为 1024。 - -```bash -VLLM_ARGS=( - --rollout-num-gpus-per-engine 64 - --vllm-gpu-memory-utilization 0.7 - --vllm-enable-expert-parallel - - # data parallelism for the attention block - --vllm-data-parallel-size 8 - - # make every dp rank has 128 concurrency - --vllm-server-concurrency 1024 -) -``` - -#### MISC_ARGS - -一些额外的 megatron 配置。注意这里配置了 megatron 的 deepep。 - -```bash -MISC_ARGS=( - ... - - # use deepep for megatron - --moe-enable-deepep - --moe-token-dispatcher-type flex -) -``` diff --git a/docs/zh/examples/glm4-9B.md b/docs/zh/examples/glm4-9B.md deleted file mode 100644 index fb53c0715..000000000 --- a/docs/zh/examples/glm4-9B.md +++ /dev/null @@ -1,280 +0,0 @@ -# 8xH100 训练 GLM4-9B - -## 环境准备 - -拉取 `inferactinc/public:vime-vllm-cu129-latest` 镜像后，用如下方式初始化镜像环境： - -```bash -cd /root/ -git clone https://github.com/THUDM/slime.git -cd slime/ -pip install -e . --no-deps -``` - -下载模型与数据： - -```bash -# hf checkpoint -hf download zai-org/GLM-Z1-9B-0414 --local-dir /root/GLM-Z1-9B-0414 - -# train data -hf download --repo-type dataset zhuzilin/dapo-math-17k \ - --local-dir /root/dapo-math-17k - -# eval data -hf download --repo-type dataset zhuzilin/aime-2024 \ - --local-dir /root/aime-2024 -``` - -将 huggingface checkpoint 转换成 megatron 可以加载的 huggingface checkpoint： - -```bash -# mcore checkpoint -cd /root/slime -source scripts/models/glm4-9B.sh -PYTHONPATH=/root/Megatron-LM python tools/convert_hf_to_torch_dist.py \ - ${MODEL_ARGS[@]} \ - --hf-checkpoint /root/GLM-Z1-9B-0414 \ - --save /root/GLM-Z1-9B-0414_torch_dist -``` - -## 执行训练 - -执行训练： - -```bash -cd /root/slime -bash script/run-glm4-9B.sh -``` - -### 参数简介 - -这里我们简单介绍一下脚本 [run-glm4-9B.sh](https://github.com/THUDM/slime/blob/main/scripts/run-glm4-9B.sh) 中的各个组成部分： - -#### MODEL_ARGS - -```bash -SCRIPT_DIR="$(cd -- "$(dirname -- "${BASH_SOURCE[0]}")" &>/dev/null && pwd)" -source "${SCRIPT_DIR}/models/glm4-9B.sh" -``` - -从 [scripts/models/glm4-9B.sh](https://github.com/THUDM/slime/blob/main/scripts/models/glm4-9B.sh) 读取模型的 config。这些 config 都是 megatron 的参数。在使用 megatron 进行训练的时候，megatron 无法从 ckpt 中读取模型 config，需要我们自行配置。我们在 [scripts/models](https://github.com/THUDM/slime/tree/main/scripts/models/) 中提供了一些样例。 - -⚠️ 注意检查模型文件中的 `--rotary-base` 等配置是否对应你当前训练模型的配置，因为同一个模型结构的不同模型可能有不同的取值。在这种情况下，你可以在导入模型参数后在脚本里进行覆盖，例如： - -```bash -source "${SCRIPT_DIR}/models/glm4-9B.sh" - -MODEL_ARGS += ( --rotary-base 10000 ) -``` - -#### CKPT_ARGS - -```bash -CKPT_ARGS=( - # vLLM 需要的 hf ckpt，我们也会从这里读 tokenizer - --hf-checkpoint /root/GLM-Z1-9B-0414 - # reference model 的 ckp - --ref-load /root/GLM-Z1-9B-0414_torch_dist - # actor 的 load dir，如果是空的，会从 `ref_load` 里面读 - --load /root/GLM-Z1-9B-0414_slime/ - --save /root/GLM-Z1-9B-0414_slime/ - --save-interval 20 -) -``` - -#### ROLLOUT_ARGS - -```bash -ROLLOUT_ARGS=( - # prompt 数据集，每行是个 json - --prompt-data /root/dapo-math-17k/dapo-math-17k.jsonl - --input-key prompt - --label-key label - # 如果 prompt 的 `input_key` 中是 openai message， - # 会进行 tokenizer.apply_chat_template(...) - --apply-chat-template - # 是否 shuffle 数据 - --rollout-shuffle - - # reward model 类型， - # slime 提供了很多类型以及用于自定义的 --custom-rm-path - --rm-type deepscaler - - # 一共要训练多少 rollout - --num-rollout 3000 - # 一个 rollout 有多少 prompt - --rollout-batch-size 32 - # 每个 prompt 采多少回复 - # 一个 rollout 会有 rollout_batch_size * n_samples_per_prompt 条 - --n-samples-per-prompt 8 - # rollout sampling param - --rollout-max-response-len 8192 - --rollout-temperature 0.8 - - # 一次 rollout 对应几个训练步 - --num-steps-per-rollout 1 - # 是否在训练时 balance data，可能对速度有好处 - --balance-data -) -``` - -#### EVAL_ARGS - -eval 的时候基本上是会继承所有 rollout 的参数，但是我们提供了一些可以 rollout 配置覆盖的参数，从而实现训练和 eval 用不同的采样策略。 - -```bash -EVAL_ARGS=( - --eval-interval 5 - --eval-prompt-data /root/aime-2024/aime-2024.jsonl - --n-samples-per-eval-prompt 16 - --eval-max-response-len 16384 - --eval-top-p 0.7 -) -``` - -#### PERF_ARGS - -一堆 megatron 的并行参数，只有 `--use-dynamic-batch-size` 与 `--max-tokens-per-gpu` 是 slime 添加的。 - -`max_tokens_per_gpu` 是指每张卡最多跑多少 token，在开启 `use_dynamic_batch_size` 之后，会尽可能将一个 batch 内部长短不一的数据拼到 `max_tokens_per_gpu`，从而组成动态的 micro batch size，如果有一条数据长度超过了 `max_tokens_per_gpu`，则自成一条，不会对数据进行截断。在开启 context parallel (CP) 时，会让 CP 张卡去上的数据去共享总长为 `CP * max_tokens_per_gpu` 的 token。 - -在开启 dynamic_batch_size，会忽略传统的 `micro_batch_size`。 - -⚠️ slime 总是会通过 data packing 的方法训练模型，并且严格保证 per sample loss 或 per token loss，也就是开启 dynamic batch size 不会对 loss 计算有影响，推荐开启。 - -```bash -PERF_ARGS=( - --tensor-model-parallel-size 2 - --sequence-parallel - --pipeline-model-parallel-size 1 - --context-parallel-size 2 - --expert-model-parallel-size 1 - --expert-tensor-parallel-size 1 - - --recompute-granularity full - --recompute-method uniform - --recompute-num-layers 1 - - # --micro-batch-size 1 - --use-dynamic-batch-size - --max-tokens-per-gpu 4608 -) -``` - -#### GRPO_ARGS - -目前 slime 这是一些 grpo 相关的参数： - -```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 -) -``` - -#### OPTIMIZER_ARGS - -```bash -OPTIMIZER_ARGS=( - --optimizer adam - --lr 1e-6 - --lr-decay-style constant - --weight-decay 0.1 - --adam-beta1 0.9 - --adam-beta2 0.98 -) -``` - -#### VLLM_ARGS - -vLLM 所需的参数，这里 `--rollout-num-gpus-per-engine` 对应 vLLM 的 `tp_size`，除此之外的 vLLM 参数均通过添加 `--vllm-` 的前缀来传给 slime。 - -```bash -VLLM_ARGS=( - --rollout-num-gpus-per-engine 2 -) -``` - -⚠️ slime 会用 vllm-router 来调度多个 vLLM 引擎。 - -### 训推一体 - -在原始的脚本中，资源配置如下： - -```bash -ray job submit ... \ - -- python3 train.py \ - --actor-num-nodes 1 \ - --actor-num-gpus-per-node 4 \ - --rollout-num-gpus 4 \ - ... -``` - -即开启训推分离，并且训练部分会使用 1 机 8 卡，推理会和训练共同使用这 8 张卡张卡。 - -如果想使用训推一体（colocate）的功能，需要加上 `--colocate` 并去掉 `--rollout-num-gpus`： - -```bash -ray job submit ... \ - -- python3 train.py \ - --actor-num-nodes 1 \ - --actor-num-gpus-per-node 8 \ - --colocate \ - ... -``` - -此时，训练和推理就会共用这 8 张卡了。 - -⚠️ 在训推一体的训练时，megatron 始终会占据一些显存，所以需要通过调整 `--vllm-gpu-memory-utilization` 来降低 vLLM 占据的显存比例。 - -### dynamic sampling - -slime 支持了更复杂的 sampling 方案，例如 [DAPO](https://dapo-sia.github.io/) 中的 dynamic sampling。如果要开启 dynamic sampling，需要配置： - -```bash - --over-sampling-batch-size ${OVER_SAMPLING_BS} \ - --dynamic-sampling-filter-path \ - slime.rollout.filter_hub.dynamic_sampling_filters.check_reward_nonzero_std \ -``` - -这里 `over_sampling_batch_size` 需要大于 ``rollout_batch_size`，例如配置为： - -```bash - --rollout-batch-size 32 \ - --n-samples-per-prompt 8 \ - --over-sampling-batch-size 64 \ -``` - -那么 sampling 会直接采样 64 条 prompt，每条 prompt 采样 8 次。因为 slime 内部进行的是异步采样，所以我们会先后获得每个 prompt 的 8 条回复。在收到回复时，会用 `dynamic_sampling_filter_path` 对应的函数进行筛选，如果通过，则留下这 8 条数据，否则则丢掉。例子中的函数是判断回答是否全对或全错： - -```python -def check_reward_nonzero_std(args, samples: list[Sample], **kwargs): - rewards = [sample.reward for sample in samples] - return torch.tensor(rewards, dtype=torch.float).std() > 0.0 -``` - -当我们收到了 32 * 8 条数据时，我们会立刻停止采样，而不会等剩余的数据采样完成。如果删除的数据超过了 32 条 prompt（剩余的小于 32 条 prompt），那么我们会再采样 64 条 prompt。 - -### partial rollout - -在进行 dynamic sampling 的过程中，会提前终止（abort）大量请求，我们可以通过配置 `--partial-rollout` 参数来将生成到一半的请求保存至 data buffer，在下一个 rollout 中取出来继续进行数据生成，从而进一步优化性能。 - -可以通过配置 `--buffer-filter-path` 来自定义如何从 buffer 中取出数据，默认的函数为： - -```python -def pop_first(args, rollout_id, buffer: list[list[Sample]], num_samples: int) -> list[list[Sample]]: - num_to_pop = min(len(buffer), num_samples) - samples = buffer[:num_to_pop] - del buffer[:num_to_pop] - return samples -``` - -即每次取出前 `num_samples` 个 prompt 对应的 `num_samples * n_samples_per_prompt` 条数据。 - -⚠️ 每条 partial rollout sample 的 `sample.metadata` 中存储了第一次进行生成的 rollout id，可以用于数据过滤。 diff --git a/docs/zh/examples/glm4.7-30B-A3B.md b/docs/zh/examples/glm4.7-30B-A3B.md deleted file mode 100644 index 0e0b005ad..000000000 --- a/docs/zh/examples/glm4.7-30B-A3B.md +++ /dev/null @@ -1,149 +0,0 @@ -# 8×H100 训练 GLM-4.7-Flash - -## 环境准备 - -搭建环境、数据与 ckpt 转换均与 Qwen3-4B 模型相同，可以参考 [示例：Qwen3-4B](qwen3-4B.md)，将文中 Qwen3-4B 的部分转换为 GLM-4.7-Flash 即可。 - -### 下载模型 - -```bash -hf download THUDM/GLM-4.7-Flash --local-dir /root/GLM-4.7-Flash -``` - -### 转换 Checkpoint - -可以用如下方法把 Hugging Face checkpoint 转化为 torch_dist 格式： - -```bash -cd /root/slime -pip install -e . --no-deps -source scripts/models/glm4.7-30B-A3B.sh -PYTHONPATH=/root/Megatron-LM/ torchrun --nproc-per-node 8 \ - tools/convert_hf_to_torch_dist.py \ - ${MODEL_ARGS[@]} \ - --hf-checkpoint /root/GLM-4.7-Flash/ \ - --save /root/GLM-4.7-Flash_torch_dist/ -``` - -## 执行训练 - -执行训练： - -```bash -cd /root/slime -bash scripts/run-glm4.7-30B-A3B-8gpus.sh -``` - -### 参数简介 - -这里我们简单介绍一下脚本 [run-glm4.7-30B-A3B-8gpus.sh](https://github.com/THUDM/slime/blob/main/scripts/run-glm4.7-30B-A3B-8gpus.sh) 中的关键部分。 - -#### MoE 配置 - -GLM-4.7-Flash 是一个 MoE（混合专家）模型，包含 64 个路由专家（top-4 激活）和 1 个共享专家。共 47 层：1 层 dense 层 + 46 层 MoE 层。 - -1. 为了支持在 8×H100 环境中运行 GLM-4.7-Flash，我们需要开启 Megatron 的 CPU Adam 以节省显存： - - ```bash - OPTIMIZER_ARGS=( - ... - --optimizer-cpu-offload - --overlap-cpu-optimizer-d2h-h2d - --use-precision-aware-optimizer - ) - ``` - -2. 开启 Megatron 支持的 MoE 优化，单机 8×H100 配置为 TP=1, EP=8： - - ```bash - PERF_ARGS=( - --tensor-model-parallel-size 1 - --pipeline-model-parallel-size 1 - --context-parallel-size 1 - --expert-model-parallel-size 8 - --expert-tensor-parallel-size 1 - ... - ) - ``` - -3. 在 vLLM 侧开启 MoE expert parallelism。GLM-4.7-Flash 是非 MLA 模型， - 这里用 attention 上 8 路 data parallel + expert 上 expert parallel： - - ```bash - VLLM_ARGS=( - --rollout-num-gpus-per-engine 8 - --vllm-gpu-memory-utilization 0.7 - --vllm-data-parallel-size 8 - --vllm-enable-expert-parallel - ... - ) - ``` - -#### MTP 投机解码（推理加速） - -GLM-4.7-Flash 包含 1 层 MTP（Multi-Token Prediction）层，可用于推理时的投机解码来加速 rollout 生成。要启用此功能，在 `VLLM_ARGS` 中添加以下配置： - -```bash -VLLM_ARGS=( - ... - # MTP 投机解码 (EAGLE) - --vllm-speculative-config '{"method":"mtp","num_speculative_tokens":3}' -) -``` - -这会让 vLLM 使用模型的 MTP 层作为 EAGLE 风格投机解码的 draft 模型。MTP 层预测多个未来 token，vLLM 并行验证它们，从而加速生成。 - -> ⚠️ **注意**：投机解码会占用额外的 GPU 显存。如果遇到 OOM 问题，可以尝试降低 `--vllm-gpu-memory-utilization` 或关闭投机解码。 - -#### MTP 训练 - -slime 也支持将 MTP 层与主模型联合训练，适用于已实现 MTP 权重转换的模型（如 MiMo、GLM-4.7）。启用时，相关参数如下： - -```bash -# 在模型配置中添加 MTP 层数 -MODEL_ARGS+=(--mtp-num-layers 1) - -# 启用 MTP 训练 -SPEC_ARGS=( - --enable-mtp-training - --mtp-loss-scaling-factor 0.2 -) -``` - -- `--mtp-num-layers 1`：告知 Megatron 从 checkpoint 中加载 MTP 层。 -- `--enable-mtp-training`：启用 MTP 层的梯度计算。不设置此标志时，MTP 层会被加载但冻结。 -- `--mtp-loss-scaling-factor 0.2`：MTP loss 相对于主策略 loss 的权重，默认为 0.2。 - -> **注意**：MTP 训练需要 MTP checkpoint bridge 正确转换 HuggingFace 和 Megatron 格式之间的权重。`GLM4MoELiteBridge`（位于 `slime_plugins/mbridge/glm4moe_lite.py`）扩展了 DeepSeek V3 bridge，实现了动态 MTP 层索引以支持 GLM-4.7-Flash 的 47 层架构。 -> -> 对于其他支持 MTP 训练的模型（如 MiMo），可参考 `scripts/run-mimo-7B-rl-eagle.sh`。 - -### 多机支持 - -对于多机训练（例如 2×8 H100），使用多机脚本： - -```bash -cd /root/slime -export BASE_DIR=/shared/path # 所有节点都可以访问的路径 -bash scripts/run-glm4.7-30B-A3B.sh -``` - -对于多机环境，需要进行如下修改： - -- 将训练模型、数据放在所有机器都可以访问到的路径上； -- 设置各台机器都可以访问到的 `MASTER_ADDR`； -- 去掉 CPU Adam 相关的配置，因为使用了 distributed optimizer，多机环境下 optimizer 的显存占比会明显下降。 -- 调整并行度：例如 TP=4, PP=2, EP=8, CP=2。 - -当总卡数并不能被 expert 总数（64）整除时，可以开启 vLLM 的 EPLB（Expert Parallelism Load Balancer），通过 `--vllm-eplb-config` 配置冗余 expert 数量。例如对于 24 卡的场景： - -```bash -VLLM_ARGS=( - --rollout-num-gpus-per-engine 24 - --vllm-gpu-memory-utilization 0.7 - --vllm-data-parallel-size 3 - --vllm-enable-expert-parallel - --vllm-enable-eplb - --vllm-eplb-config '{"num_redundant_experts": 16}' -) -``` diff --git a/docs/zh/examples/glm4.7-355B-A32B.md b/docs/zh/examples/glm4.7-355B-A32B.md deleted file mode 100644 index 56519fca2..000000000 --- a/docs/zh/examples/glm4.7-355B-A32B.md +++ /dev/null @@ -1,183 +0,0 @@ -# 64xH100 训练 GLM-4.7 - -## 环境准备 - -搭建环境与下载数据的方法与 Qwen3-4B 模型相同，可以参考 [示例：Qwen3-4B](qwen3-4B.md)，将文中 Qwen3-4B 的部分替换为 GLM-4.7 即可。 - -### 前置条件 - -GLM-4.7 使用 slime 标准 Docker 环境即可。多机启动前，请确保所有机器都能访问同一个 `$BASE_DIR` 路径，并在启动 Ray worker 前先取消代理： - -```bash -unset http_proxy https_proxy HTTP_PROXY HTTPS_PROXY -``` - -### 下载模型 - -```bash -hf download zai-org/GLM-4.7 --local-dir $BASE_DIR/GLM-4.7-355B-A32B -``` - -### 转换 Checkpoint - -可以用如下方法把 Hugging Face checkpoint 转换为 torch_dist 格式（2 机 x 8 卡）： - -```bash -cd /root/slime -pip install -e . --no-deps -source scripts/models/glm4.5-355B-A32B.sh -PYTHONPATH=/root/Megatron-LM/ torchrun \ - --nproc-per-node 8 \ - --master-addr ${MASTER_ADDR} --master-port 12345 \ - --nnodes=2 --node-rank ${NODE_RANK} \ - tools/convert_hf_to_torch_dist.py \ - ${MODEL_ARGS[@]} \ - --hf-checkpoint $BASE_DIR/GLM-4.7-355B-A32B/ \ - --save $BASE_DIR/GLM-4.7-355B-A32B_torch_dist/ -``` - -其中 `MASTER_ADDR` 是 node0 的 IP，`NODE_RANK` 表示当前机器的编号，配置方式与普通多机 `torchrun` 一致。 - -## 执行训练 - -从 node0 执行训练脚本： - -```bash -cd /root/slime -export BASE_DIR=/shared/path # 所有节点都能访问的共享路径 -bash scripts/run-glm4.7-355B-A32B.sh -``` - -### 参数简介 - -这里我们简单介绍一下 [run-glm4.7-355B-A32B.sh](https://github.com/THUDM/slime/blob/main/scripts/run-glm4.7-355B-A32B.sh) 中的关键部分。 - -#### MoE 配置 - -GLM-4.7 是一个 MoE（混合专家）模型，包含 160 个路由专家（top-8 激活）和共享专家。模型共 92 层：3 层 dense + 89 层 MoE。 - -1. 为了支持在 64xH100 环境中运行 GLM-4.7，我们开启 Megatron 的 CPU Adam 来节省显存： - - ```bash - OPTIMIZER_ARGS=( - ... - --optimizer-cpu-offload - --overlap-cpu-optimizer-d2h-h2d - --use-precision-aware-optimizer - ) - ``` - -2. 在 Megatron 中开启 MoE 优化。当前 64xH100 示例使用 TP=8、PP=4、CP=2、EP=16： - - ```bash - PERF_ARGS=( - --tensor-model-parallel-size 8 - --sequence-parallel - --pipeline-model-parallel-size 4 - --context-parallel-size 2 - --expert-model-parallel-size 16 - --expert-tensor-parallel-size 1 - ... - --use-dynamic-batch-size - --max-tokens-per-gpu 16384 - ) - ``` - -3. 在 vLLM 侧开启 MoE expert parallelism。GLM-4.7 是非 MLA 模型， - 这里用 attention 上 4 路 data parallel + expert 上 expert parallel - （EP size 由 `tensor_parallel_size × data_parallel_size` 自动推导）： - - ```bash - VLLM_ARGS=( - --rollout-num-gpus-per-engine 32 - --vllm-gpu-memory-utilization 0.7 - --vllm-data-parallel-size 4 - --vllm-enable-expert-parallel - ... - ) - ``` - -#### MTP 投机解码（推理加速） - -GLM-4.7 包含 MTP（Multi-Token Prediction）层，可以在推理阶段用于投机解码，加速 rollout 生成。启用方法是在 `VLLM_ARGS` 中加入： - -```bash -VLLM_ARGS=( - ... - # MTP 投机解码 (EAGLE) - --vllm-speculative-config '{"method":"mtp","num_speculative_tokens":3}' -) -``` - -这样 vLLM 就会使用模型自带的 MTP 层作为 EAGLE 风格投机解码的 draft model。 - -> ⚠️ **注意**：投机解码会额外占用 GPU 显存。如果遇到 OOM，可以尝试降低 `--vllm-gpu-memory-utilization` 或暂时关闭投机解码。 - -#### MTP 训练 - -slime 也支持在 GLM-4.7 上将 MTP 层与主模型联合训练。启用时，相关参数如下： - -```bash -# 在模型配置中添加 MTP 层数 -MODEL_ARGS+=(--mtp-num-layers 1) - -# 启用 MTP 训练 -MTP_ARGS=( - --enable-mtp-training - --mtp-loss-scaling-factor 0.2 -) -``` - -- `--mtp-num-layers 1`：告知 Megatron 从 checkpoint 中加载 MTP 层。 -- `--enable-mtp-training`：启用 MTP 层的梯度计算；不设置时 MTP 层会被加载但保持冻结。 -- `--mtp-loss-scaling-factor 0.2`：MTP loss 相对主策略 loss 的权重，默认值为 0.2。 - -> **注意**：GLM-4.7 的 MTP 训练依赖 `GLM4MoEBridge`（位于 `slime_plugins/mbridge/glm4moe.py`）在 HuggingFace 与 Megatron 格式之间正确映射普通层和 MTP 层权重。 - -#### 多机支持 - -这个示例本身就是多机训练配置。启动前请确认： - -- 模型权重和数据集放在所有节点都能访问到的路径； -- `MASTER_ADDR` 设置为所有节点都能访问到的地址； -- 在启动 Ray worker 前先取消代理； -- 提供一个 `HOSTFILE` 列出 worker IP（每行一个），并在启动前 `export HOSTFILE=/path/to/hostfile`； -- 并行度需要成套调整。默认示例使用 TP=8、PP=4、EP=16、CP=2，rollout 侧则使用 32 张卡 / engine + DP=4 + expert parallel。 - -如果 rollout GPU 数与 expert 数（160）之间不能整除，可以开启 vLLM 的 EPLB（Expert Parallelism Load Balancer），通过 `--vllm-eplb-config` 配置冗余 expert，例如： - -```bash ---vllm-enable-eplb ---vllm-eplb-config '{"num_redundant_experts": 16}' -``` - -## FP8 Rollout - -开源版 GLM-4.7 的 FP8 checkpoint 使用的是 per-channel 量化，目前无法在 vLLM 中直接启用 DeepEP。可以利用 slime 自带工具将其转换为 128x128 的 per-block FP8 checkpoint： - -```bash -cd /root/slime -python tools/convert_hf_to_fp8.py \ - --model-dir $BASE_DIR/GLM-4.7-355B-A32B/ \ - --save-dir $BASE_DIR/GLM-4.7-355B-A32B-FP8/ \ - --strategy block --block-size 128 128 \ - --max-workers 4 -``` - -随后把 `--hf-checkpoint` 改成 `$BASE_DIR/GLM-4.7-355B-A32B-FP8/` 即可开启 FP8 rollout。 - -一个可参考的 FP8 `VLLM_ARGS` 配置如下： - -```bash -VLLM_ARGS=( - --rollout-num-gpus-per-engine 32 - --vllm-gpu-memory-utilization 0.7 - --vllm-data-parallel-size 32 - --vllm-enable-expert-parallel - --vllm-cudagraph-capture-sizes 1 2 4 8 $(seq 16 8 128) - - --vllm-speculative-config '{"method":"mtp","num_speculative_tokens":3}' - - --vllm-all2all-backend deepep_high_throughput -) -``` diff --git a/docs/zh/examples/qwen3-4b-base-openhermes.md b/docs/zh/examples/qwen3-4b-base-openhermes.md deleted file mode 100644 index 3de99a22b..000000000 --- a/docs/zh/examples/qwen3-4b-base-openhermes.md +++ /dev/null @@ -1,85 +0,0 @@ -# SFT Qwen3-4B-Base - -## 环境准备 - -首先需要我们仿照 [示例：Qwen3-4B 模型](qwen3-4B.md) 创建镜像环境与转换 `Qwen3-4B-Base` 模型。 - -之后，我们处理 sft 数据。这里我们以经典的 [OpenHermes-2.5](https://huggingface.co/datasets/teknium/OpenHermes-2.5) 为例，首先把数据处理成适合 slime 加载的格式，可以用如下的脚本进行处理，增加一个符合 openai message 格式的列，并保存在 `/root/openhermes2_5.parquet`。 - -```python -from datasets import load_dataset - -ds = load_dataset("teknium/OpenHermes-2.5")["train"] - -def convert(sample): - conversations = sample["conversations"] - - def convert_role(role): - if role == "human": - return "user" - elif role == "gpt": - return "assistant" - elif role == "system": - return "system" - else: - raise ValueError(f"Unknown role: {role}") - - messages = [ - { - "role": convert_role(turn["from"]), - "content": turn["value"], - } - for turn in conversations - ] - - return {"messages": messages} - -ds = ds.map(convert) -ds.to_parquet("/root/openhermes2_5.parquet") -``` - -## 执行训练 - -执行训练： - -```bash -cd /root/slime -bash script/run-qwen3-4B-base-sft.sh -``` - -### 参数简介 - -可以将 [run-qwen3-4B-base-sft.sh](https://github.com/THUDM/slime/blob/main/scripts/run-qwen3-4B-base-sft.sh) 与 [run-qwen3-4B.sh](https://github.com/THUDM/slime/blob/main/scripts/run-qwen3-4B.sh) 进行对比。会发现除了我们将模型由 instruct 模型换为了 base 模型之外，主要进行了如下的几个调整： - -1. 移除了 `VLLM_ARGS` 和 `GRPO_ARGS`。这是因为 sft 的过程中不需要启动 vLLM 或者做 grpo 相关的配置； - -2. 将 `ROLLOUT_ARGS` 改名为了 `SFT_ARGS`，并配置为： - - ```bash - SFT_ARGS=( - --rollout-function-path slime.rollout.sft_rollout.generate_rollout - --prompt-data /root/openhermes2_5.parquet - --input-key messages - --rollout-shuffle - --num-epoch 3 - --rollout-batch-size 128 - --global-batch-size 128 - - --loss-type sft_loss - --calculate-per-token-loss - --disable-compute-advantages-and-returns - --debug-train-only - ) - ``` - - slime 中的 sft 实际上是复用了 slime 的 custom rollout 功能，通过 `--rollout-function-path` 将数据生成部分从使用 vLLM 的 RL rollout，切换成了从文件中读取数据的 sft 版本，即 `slime.rollout.sft_rollout.generate_rollout`。 - - 对于 sft 来说，建议将 `rollout_batch_size` 与 `global_batch_size` 设置成相同的，并不要配置 `n_samples_per_prompt`，这样相当于是读一个 batch 就训一个 batch。 - - slime 还支持不同的 loss 类型，我们就是通过 `--loss-type sft_loss` 配置上 sft loss 的。 - - 至于 `--calculate-per-token-loss`，这是因为 slime 默认是以 GRPO 的 per sample mean 进行计算的，而一般 sft 训练都是按一个 batch 的所有不被 mask 的 token 取平均，所以建议配置上。 - - 最后 `--disable-compute-advantages-and-returns` 表示 sft 的过程中不需要预先计算 log prob，`--debug-train-only` 表示不需要初始化 vLLM。 - -3. 使用了 `train_async.py` 而不是 `train.py`。这是为了利用异步训练的流程，来实现数据 prefetch。 diff --git a/docs/zh/examples/qwen3-next-80B-A3B.md b/docs/zh/examples/qwen3-next-80B-A3B.md deleted file mode 100644 index 47997a63a..000000000 --- a/docs/zh/examples/qwen3-next-80B-A3B.md +++ /dev/null @@ -1,98 +0,0 @@ -# 8xH100 训练 Qwen3-30B-A3B - -## 环境准备 - -搭建环境、下载模型、数据与 ckpt 转换均与 Qwen3-4B 模型相同，可以参考 [示例：Qwen3-4B](./qwen3-4B.md)，将文中 Qwen3-4B 的部分转换为 -Qwen3-next-80B-A3B-Instruct 即可。 - -可以用如下完整方法把 huggingface checkpoint 转化为 torch_dist 格式： - -```bash -export BASE_FOLDER=./models/ -# 下载模型权重 (Qwen3-Next-80B-A3B-Thinking) -hf download Qwen/Qwen3-Next-80B-A3B-Thinking --local-dir ${BASE_FOLDER}/Qwen3-Next-80B-A3B-Thinking -``` - -```shell -cd slime/ -pip install -e . --no-deps - -# (for acceleration) -cd .. # and find a proper folder -git clone https://github.com/fla-org/flash-linear-attention -cd flash-linear-attention -git checkout 9714c595 -pip install -e . --no-deps - -wget https://github.com/Dao-AILab/causal-conv1d/releases/download/v1.5.4/causal_conv1d-1.5.4+cu12torch2.8cxx11abiTRUE-cp312-cp312-linux_x86_64.whl -pip install ./causal_conv1d-1.5.4+cu12torch2.8cxx11abiTRUE-cp312-cp312-linux_x86_64.whl -``` - -## [Optional] Fix a bug in triton compilation on Blackwell (sm100) - -see discussion here https://github.com/triton-lang/triton/issues/8695 -and https://github.com/fla-org/flash-linear-attention/issues/638 - -We need to apply a patch to fix the bug. -Go to the flash-linear-attention folder you just installed, and apply the following patch: - -```diff -diff --git a/fla/ops/gated_delta_rule/wy_fast.py b/fla/ops/gated_delta_rule/wy_fast.py -index c5119dcf..838f5e4e 100644 ---- a/fla/ops/gated_delta_rule/wy_fast.py -+++ b/fla/ops/gated_delta_rule/wy_fast.py -@@ -198,7 +198,14 @@ def prepare_wy_repr_bwd_kernel( - b_A += tl.dot(b_kb, tl.trans(b_k)) - b_dkb = tl.dot(b_dA, b_k) - b_db += tl.sum(b_dkb * b_k, 1) -- b_dk += tl.dot(tl.trans(b_dA), b_kb) -+ b_dk += tl.inline_asm_elementwise( -+ asm="mov.f32 $0, $1;", -+ constraints="=r,r", -+ args=[tl.dot(tl.trans(b_dA), b_kb)], -+ dtype=tl.float32, -+ is_pure=True, -+ pack=1, -+ ) - b_dk += b_dkb * b_b[:, None] - tl.store(p_dk, b_dk.to(p_dk.dtype.element_ty), boundary_check=(0, 1)) - tl.store(p_db, b_db.to(p_db.dtype.element_ty), boundary_check=(0,)) - -``` - -save it as `patch.diff` (Please remember to copy the last empty line to the file!) and do `git apply patch.diff` - -## 执行训练 (Megatron) - -**当前暂不支持Blackwell** - -转换模型权重： - -```bash -source scripts/models/qwen3-next-80B-A3B.sh -PYTHONPATH=/root/Megatron-LM/ torchrun --nproc-per-node 8 \ - tools/convert_hf_to_torch_dist.py \ - ${MODEL_ARGS[@]} \ - --hf-checkpoint /root/Qwen3-Next-80B-A3B-Thinking/ \ - --save /root/Qwen3-Next-80B-A3B-Thinking_torch_dist/ -``` - -单机8卡 - -```bash -cd /root/slime -export BASE_FOLDER=/root -export MASTER_ADDR=127.0.0.1 -bash scripts/run-qwen3-next-80B-A3B-8gpus.sh -``` -如果显存不够，考虑disable `--accumulate-allreduce-grads-in-fp32`，enable `--grad-reduce-in-bf16` - - -多机（4x8） - -```bash -cd /root/slime -export BASE_FOLDER=/root -export MASTER_ADDR=your_master_addr -bash scripts/run-qwen3-next-80B-A3B.sh -``` diff --git a/docs/zh/get_started/quick_start.md b/docs/zh/get_started/quick_start.md index 0dff33a4c..e8bc4458c 100644 --- a/docs/zh/get_started/quick_start.md +++ b/docs/zh/get_started/quick_start.md @@ -52,7 +52,7 @@ pip install -e . --no-deps ```bash # 下载模型权重 (GLM-Z1-9B) -hf download zai-org/GLM-Z1-9B-0414 --local-dir /root/GLM-Z1-9B-0414 +hf download zai-org/Qwen3-4B --local-dir /root/Qwen3-4B # 下载训练数据集 (dapo-math-17k) hf download --repo-type dataset zhuzilin/dapo-math-17k \ @@ -69,11 +69,11 @@ hf download --repo-type dataset zhuzilin/aime-2024 \ 当使用 Megatron 作为训练后端时，需要先将 Hugging Face 格式的模型权重转换为 Megatron `torch_dist` 格式。 -首先，加载目标模型的配置文件。`slime/scripts/models` 目录下包含了支持模型的配置文件。需要 `source` 对应模型的脚本，将配置参数加载到当前环境中。此处我们以 GLM4-9B 模型为例子，对于 Qwen3-4B、Qwen3.5、Qwen3.6、GLM-4.7-Flash、Qwen3-30B-A3B，是类似的。 +首先，加载目标模型的配置文件。`slime/scripts/models` 目录下包含了支持模型的配置文件。需要 `source` 对应模型的脚本，将配置参数加载到当前环境中。此处我们以 Qwen3-4B 模型为例子。 ```bash cd /root/slime -source scripts/models/glm4-9B.sh +source scripts/models/qwen3-4B.sh ``` 接下来，运行转换脚本。请注意以下参数： @@ -83,12 +83,11 @@ source scripts/models/glm4-9B.sh ```bash PYTHONPATH=/root/Megatron-LM python tools/convert_hf_to_torch_dist.py \ ${MODEL_ARGS[@]} \ - --hf-checkpoint /root/GLM-Z1-9B-0414 \ - --save /root/GLM-Z1-9B-0414_torch_dist + --hf-checkpoint /root/Qwen3-4B \ + --save /root/Qwen3-4B_torch_dist ``` 对于更大的模型，可以使用 `torchrun` 来启动转换脚本，从而使用多张 GPU 甚至多机进行权重转换。 -注意：kimi-k2模型权重转换时，需打开模型路径中的config.json，将"model_type": "kimi_k2"修改为"model_type": "deepseek_v3"。 ### Megatron 格式 转换为 Hugging Face 格式 @@ -97,8 +96,8 @@ PYTHONPATH=/root/Megatron-LM python tools/convert_hf_to_torch_dist.py \ ```bash PYTHONPATH=/root/Megatron-LM python tools/convert_torch_dist_to_hf.py \ --input-dir /path/to/torch_dist_ckpt/iter_xxx/ \ - --output-dir /root/GLM-Z1-9B-0414-iter_xxx \ - --origin-hf-dir /root/GLM-Z1-9B-0414 + --output-dir /root/Qwen3-4B-iter_xxx \ + --origin-hf-dir /root/Qwen3-4B ``` 由于 Megatron 会对 embedding 做 padding，可能会出现转换出来的权重的 embedding 形状不匹配的问题。这时需要在转换时设置 `--vocab-size`。 @@ -109,24 +108,24 @@ PYTHONPATH=/root/Megatron-LM python tools/convert_torch_dist_to_hf.py \ ```bash cd /root/slime -bash scripts/run-glm4-9B.sh +bash scripts/run-qwen3-4B.sh ``` -我们还是以 run-glm4-9B.sh 脚本为例，简单分析主要参数的作用。 +我们还是以 run-qwen3-4B.sh 脚本为例，简单分析主要参数的作用。 ### MODEL_ARGS: 模型配置参数 ```bash SCRIPT_DIR="$(cd -- "$(dirname -- "${BASH_SOURCE[0]}")" &>/dev/null && pwd)" -source "${SCRIPT_DIR}/models/glm4-9B.sh" +source "${SCRIPT_DIR}/models/qwen3-4B.sh" ``` -此部分通过 `source` 命令从 `scripts/models/glm4-9B.sh` 文件中加载模型配置。这些配置均为 Megatron 所需的超参数。由于 Megatron 无法直接从检查点（checkpoint）中读取模型配置，因此需要手动指定。我们在 `scripts/models/` 目录下提供了一些常用模型的配置示例。 +此部分通过 `source` 命令从 `scripts/models/qwen3-4B.sh` 文件中加载模型配置。这些配置均为 Megatron 所需的超参数。由于 Megatron 无法直接从检查点（checkpoint）中读取模型配置，因此需要手动指定。我们在 `scripts/models/` 目录下提供了一些常用模型的配置示例。 > ⚠️ **注意**： > 请务必检查模型配置文件中的参数（如 `--rotary-base`）是否与您当前使用的模型完全匹配。同一模型结构的不同版本可能使用不同的配置值。如果需要修改，您可以在 `source` 之后直接覆盖，例如： > ```bash -> source "${SCRIPT_DIR}/models/glm4-9B.sh" +> source "${SCRIPT_DIR}/models/qwen3-4B.sh" > MODEL_ARGS+=(--rotary-base 10000) > ``` @@ -135,13 +134,13 @@ source "${SCRIPT_DIR}/models/glm4-9B.sh" ```bash CKPT_ARGS=( # 用于加载 tokenizer 等其他信息，实际上不会使用 hf 路径中的模型权重参数 - --hf-checkpoint /root/GLM-Z1-9B-0414 + --hf-checkpoint /root/Qwen3-4B # 参考模型 (Reference Model) 的 Megatron 格式检查点 - --ref-load /root/GLM-Z1-9B-0414_torch_dist + --ref-load /root/Qwen3-4B_torch_dist # Actor 模型的加载路径。若为空或不存在有效的checkpoint，则从 --ref-load 加载 - --load /root/GLM-Z1-9B-0414_slime/ + --load /root/Qwen3-4B_slime/ # 训练过程中模型的保存路径 - --save /root/GLM-Z1-9B-0414_slime/ + --save /root/Qwen3-4B_slime/ # 模型保存间隔（步数） --save-interval 20 ) @@ -574,8 +573,6 @@ ray job submit --address="http://127.0.0.1:8265" \ --...（其他 Megatron/vLLM/slime 参数） ``` -slime 针对大规模混合专家（MoE）模型的分布式训练进行了深度优化。我们提供了一些端到端的训练案例以供参考： +slime 针对大规模混合专家（MoE）模型的分布式训练进行了深度优化。我们提供了一个端到端的训练案例以供参考： -- [示例：8xH100 训练 GLM-4.7-Flash](../examples/glm4.7-30B-A3B.md) -- [示例：64xH100 训练 GLM-4.7](../examples/glm4.7-355B-A32B.md) -- [示例：128xH100 训练 DeepSeek-R1](../examples/deepseek-r1.md) +- [示例：8xH100 训练 Qwen3-30B-A3B](../examples/qwen3-30B-A3B.md) diff --git a/docs/zh/index.rst b/docs/zh/index.rst index 98b960608..be92352ad 100644 --- a/docs/zh/index.rst +++ b/docs/zh/index.rst @@ -26,24 +26,18 @@ slime 是 GLM-4.7、GLM-4.6、GLM-4.5 背后的 RL 训练框架。除此之外 :caption: Dense examples/qwen3-4B.md - examples/glm4-9B.md .. toctree:: :maxdepth: 1 :caption: MoE - examples/glm4.7-30B-A3B.md examples/qwen3-30B-A3B.md - examples/glm4.7-355B-A32B.md - examples/deepseek-r1.md .. toctree:: :maxdepth: 1 :caption: 高级特性 - advanced/on-policy-distillation.md advanced/speculative-decoding.md - advanced/low-precision.md advanced/reproducibility.md advanced/fault-tolerance.md advanced/pd-disaggregation.md @@ -55,7 +49,6 @@ slime 是 GLM-4.7、GLM-4.6、GLM-4.5 背后的 RL 训练框架。除此之外 :maxdepth: 1 :caption: 其他用法 - examples/qwen3-4b-base-openhermes.md _examples_synced/search-r1/README.md _examples_synced/fully_async/README.md _examples_synced/retool/README.md @@ -69,4 +62,3 @@ slime 是 GLM-4.7、GLM-4.6、GLM-4.5 背后的 RL 训练框架。除此之外 developer_guide/debug.md developer_guide/trace.md developer_guide/profiling.md - diff --git a/examples/eval_multi_task/README.md b/examples/eval_multi_task/README.md deleted file mode 100644 index 0bf61ec72..000000000 --- a/examples/eval_multi_task/README.md +++ /dev/null @@ -1,12 +0,0 @@ -# Multi-Task Evaluation Example - -## Configuring `multi_task.yaml` -- `eval.defaults` defines inference parameters shared by every dataset entry. Override them inside an individual dataset block if needed. -- `eval.datasets` enumerates the datasets to evaluate. Each entry should specify: - - `name`: a short identifier that appears in logs and dashboards. - - `path`: the path to the dataset JSONL file. - - `rm_type`: which reward function to use for scoring. - - `n_samples_per_eval_prompt`: how many candidate completions to generate per prompt. - -## IFBench Notes -- When `ifbench` is used, `slime/rollout/rm_hub/ifbench.py` will automatically prepares the scoring environment, so no additional manual setup is required beyond providing the dataset path. diff --git a/examples/eval_multi_task/multi_task.sh b/examples/eval_multi_task/multi_task.sh deleted file mode 100644 index f15b29650..000000000 --- a/examples/eval_multi_task/multi_task.sh +++ /dev/null @@ -1,149 +0,0 @@ -#!/bin/bash - -# for rerun the task -pkill -9 -f "vllm serve" -sleep 3 -ray stop --force -pkill -9 ray -pkill -9 python -sleep 3 -pkill -9 ray -pkill -9 python - -set -ex - -# will prevent ray from buffering stdout/stderr -export PYTHONBUFFERED=16 - -NVLINK_COUNT=$(nvidia-smi topo -m 2>/dev/null | grep -o 'NV[0-9][0-9]*' | wc -l) -if [ "$NVLINK_COUNT" -gt 0 ]; then - HAS_NVLINK=1 -else - HAS_NVLINK=0 -fi -echo "HAS_NVLINK: $HAS_NVLINK (detected $NVLINK_COUNT NVLink references)" - -SCRIPT_DIR="$(cd -- "$(dirname -- "${BASH_SOURCE[0]}")" &>/dev/null && pwd)" -REPO_ROOT="$(cd -- "${SCRIPT_DIR}/../.." &>/dev/null && pwd)" -source "${REPO_ROOT}/scripts/models/qwen3-4B.sh" -EVAL_CONFIG_PATH="${REPO_ROOT}/examples/eval_multi_task/multi_task.yaml" - -CKPT_ARGS=( - --hf-checkpoint /root/Qwen3-4B - #--hf-checkpoint /root/Qwen3-4B-FP8 - --ref-load /root/Qwen3-4B_torch_dist - --load /root/Qwen3-4B_slime/ - --save /root/Qwen3-4B_slime/ - --save-interval 20 -) - -ROLLOUT_ARGS=( - --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_ARGS=( - --eval-interval 20 - --eval-config "${EVAL_CONFIG_PATH}" -) - -PERF_ARGS=( - --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 - - # --micro-batch-size 1 - --use-dynamic-batch-size - --max-tokens-per-gpu 9216 -) - -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 -) - -OPTIMIZER_ARGS=( - --optimizer adam - --lr 1e-6 - --lr-decay-style constant - --weight-decay 0.1 - --adam-beta1 0.9 - --adam-beta2 0.98 -) - -WANDB_ARGS=( - --use-wandb - --wandb-project eval - --wandb-group multi_task - --wandb-key ${WANDB_KEY} -) - -VLLM_ARGS=( - --rollout-num-gpus-per-engine 2 - --vllm-gpu-memory-utilization 0.7 -) - -MISC_ARGS=( - # default dropout in megatron is 0.1 - --attention-dropout 0.0 - --hidden-dropout 0.0 - # should be good for model performance - --accumulate-allreduce-grads-in-fp32 - --attention-softmax-in-fp32 - # need to comment this when using model with MLA - --attention-backend flash -) - -# launch the master node of ray in container -export MASTER_ADDR=${MASTER_ADDR:-"127.0.0.1"} -ray start --head --node-ip-address ${MASTER_ADDR} --num-gpus 8 --disable-usage-stats --dashboard-host=0.0.0.0 --dashboard-port=8265 - -# Build the runtime environment JSON with proper variable substitution -RUNTIME_ENV_JSON="{ - \"env_vars\": { - \"PYTHONPATH\": \"/root/Megatron-LM/\", - \"CUDA_DEVICE_MAX_CONNECTIONS\": \"1\", - \"NCCL_NVLS_ENABLE\": \"${HAS_NVLINK}\" - } -}" - -ray job submit --address="http://127.0.0.1:8265" \ - --runtime-env-json="${RUNTIME_ENV_JSON}" \ - -- python3 "${REPO_ROOT}/train.py" \ - --actor-num-nodes 1 \ - --actor-num-gpus-per-node 8 \ - --colocate \ - ${MODEL_ARGS[@]} \ - ${CKPT_ARGS[@]} \ - ${ROLLOUT_ARGS[@]} \ - ${OPTIMIZER_ARGS[@]} \ - ${GRPO_ARGS[@]} \ - ${WANDB_ARGS[@]} \ - ${PERF_ARGS[@]} \ - ${EVAL_ARGS[@]} \ - ${VLLM_ARGS[@]} \ - ${MISC_ARGS[@]} diff --git a/examples/eval_multi_task/multi_task.yaml b/examples/eval_multi_task/multi_task.yaml deleted file mode 100644 index bad2d6141..000000000 --- a/examples/eval_multi_task/multi_task.yaml +++ /dev/null @@ -1,17 +0,0 @@ -eval: - defaults: - max_response_len: 16384 - top_p: 0.7 - datasets: - - name: aime - path: /root/aime-2024/aime-2024.jsonl - rm_type: deepscaler - n_samples_per_eval_prompt: 16 - - name: gpqa # hf download --repo-type dataset zyzshishui0627/gpqa_diamond --local-dir /root/gpqa - path: /root/gpqa/gpqa_eval.jsonl - rm_type: gpqa - n_samples_per_eval_prompt: 2 - - name: ifbench # hf download --repo-type dataset zyzshishui0627/IFBench --local-dir /root/ifbench - path: /root/ifbench/IFBench_eval.jsonl - rm_type: ifbench - n_samples_per_eval_prompt: 1 diff --git a/examples/eval_multi_task/requirements_ifbench.txt b/examples/eval_multi_task/requirements_ifbench.txt deleted file mode 100644 index 78f13fac4..000000000 --- a/examples/eval_multi_task/requirements_ifbench.txt +++ /dev/null @@ -1,6 +0,0 @@ -emoji -immutabledict -nltk -numpy==1.26.4 -spacy==3.7.4 -syllapy diff --git a/examples/on_policy_distillation/README.md b/examples/on_policy_distillation/README.md deleted file mode 100644 index 9f523ec02..000000000 --- a/examples/on_policy_distillation/README.md +++ /dev/null @@ -1,127 +0,0 @@ -# On-Policy Distillation Example - -This example shows how to run **on-policy distillation (OPD)** using slime. A small student (Qwen3-8B) is aligned to imitate a larger teacher (Qwen3-32B) by training only on the student's own rollouts and matching the teacher's token-level log-probabilities. - -## Key Features - -- **OPD is orthogonal to advantage estimators**: OPD works as an additive KL penalty on top of any advantage estimator (GRPO, PPO, REINFORCE++, etc.), not as a separate estimator. -- **Two teacher modes**: - - **vllm**: Teacher runs on an external vLLM server, teacher log-probs are obtained during rollout via `/v1/completions` with `prompt_logprobs`. - - **megatron**: Teacher is loaded directly into Megatron via `--opd-teacher-load`, teacher log-probs are computed during training forward pass. - -## Key Arguments - -| Argument | Description | -|----------|-------------| -| `--use-opd` | Enable on-policy distillation. Required flag to use OPD. | -| `--opd-type` | Type of OPD: `vllm` or `megatron`. Required when `--use-opd` is set. | -| `--opd-kl-coef` | OPD KL penalty coefficient (default: 1.0). | -| `--opd-teacher-load` | Path to teacher checkpoint. **Required** when `--opd-type=megatron`, **must not be set** when `--opd-type=vllm`. | -| `--opd-teacher-ckpt-step` | Optional checkpoint step for teacher model. | - -## Mode Comparison - -| Mode | Teacher Location | When to use | -|------|------------------|-------------| -| `vllm` | External vLLM server | Teacher has different architecture or larger than GPU memory | -| `megatron` | Loaded into Megatron training | Teacher has same architecture as policy/ref model | - -## Components - -- `slime/rollout/on_policy_distillation.py` implements (for vLLM mode): - - `reward_func` calls the teacher server (via `args.rm_url`, an OpenAI-compatible `/v1/completions` endpoint) with every sample to obtain prompt log-probs. - - `post_process_rewards` trims the teacher logprobs to the generated response span and writes the tensors back to each `Sample` to compute advantages. -- `run-qwen3-8B-opd.sh` launches a vLLM teacher server, then submits a Ray job that runs `train.py`. -- `run-qwen3-8B-opd-megatron.sh` uses Megatron-loaded teacher model (no external server needed). - -## Running the example - -### Using vLLM Teacher (External Server) - -1. Download or prepare the required checkpoints and data. -```bash -hf download Qwen/Qwen3-32B --local-dir /root/Qwen3-32B -hf download Qwen/Qwen3-8B --local-dir /root/Qwen3-8B -hf download --repo-type dataset zhuzilin/dapo-math-17k --local-dir /root/dapo-math-17k -``` - -2. Run the hf to mcore for student model conversion: -```bash -cd /root/slime -source scripts/models/qwen3-8B.sh - -PYTHONPATH=/root/Megatron-LM python tools/convert_hf_to_torch_dist.py \ - ${MODEL_ARGS[@]} \ - --hf-checkpoint /root/Qwen3-8B \ - --save /root/Qwen3-8B_torch_dist -``` - -3. Run on-policy distillation: -```bash -bash examples/on_policy_distillation/run-qwen3-8B-opd.sh -``` - -### Using Megatron Teacher (No External Server) - -1. Prepare student checkpoint (same as above). - -2. **IMPORTANT**: Convert your teacher model to Megatron format (change the path to your actual teacher): -```bash -# This example uses the same model as both student and teacher (for demonstration only) -# In practice, use a different (stronger) model as the teacher! -cd /root/slime -source scripts/models/qwen3-8B.sh # Or your teacher model config - -PYTHONPATH=/root/Megatron-LM python tools/convert_hf_to_torch_dist.py \ - ${MODEL_ARGS[@]} \ - --hf-checkpoint /root/YourTeacherModel \ - --save /root/YourTeacherModel_torch_dist -``` - -3. Edit `run-qwen3-8B-opd-megatron.sh` to update paths: - - Change `--opd-teacher-load` to your teacher model path - - Adjust `--opd-kl-coef` based on your task - -4. Run: -```bash -bash examples/on_policy_distillation/run-qwen3-8B-opd-megatron.sh -``` - - -# Preliminary Results -Using Qwen3-8B-Base model sfted on part of the [OpenThoughts3-1.2M](https://huggingface.co/datasets/open-thoughts/OpenThoughts3-1.2M) dataset, we performed on-policy distillation with a Qwen3-32B teacher on the remaining data. Evaluation on Math500 shows: - -| | Pass@1 | -|-----------------------------------------------|--------| -| Qwen3-8B-Base + SFT | 76% | -| Qwen3-8B-Base + SFT + On-Policy Distillation | 94% | - - - - - -# FAQ -1. **Why are there two OPD modes?** - - `vllm` mode: The teacher runs on an independent vLLM server. This is useful when the teacher has a different architecture or is too large to load together with the policy model. - - `megatron` mode: The teacher is loaded into Megatron using the same parameter loading mechanism as the reference model. This requires the teacher to have the same architecture as the policy model. - -2. **How do I use Megatron-based teacher instead of vLLM server?** - Replace your OPD arguments: - ```bash - # Instead of: - --use-opd --opd-type vllm --opd-kl-coef 1.0 - # Use: - --use-opd --opd-type megatron --opd-kl-coef 1.0 --opd-teacher-load /path/to/teacher_checkpoint - ``` - -3. **What happens if I set wrong arguments?** - The system will raise clear errors: - - `--use-opd` without `--opd-type`: Error asking you to specify type - - `--opd-type megatron` without `--opd-teacher-load`: Error asking for teacher checkpoint - - `--opd-type vllm` with `--opd-teacher-load`: Error indicating conflict - - -# References -1. https://thinkingmachines.ai/blog/on-policy-distillation/ -2. https://arxiv.org/abs/2306.13649 -3. https://arxiv.org/abs/2306.08543 \ No newline at end of file diff --git a/examples/on_policy_distillation/run-qwen3-8B-opd-megatron.sh b/examples/on_policy_distillation/run-qwen3-8B-opd-megatron.sh deleted file mode 100644 index 248b63480..000000000 --- a/examples/on_policy_distillation/run-qwen3-8B-opd-megatron.sh +++ /dev/null @@ -1,166 +0,0 @@ -#!/bin/bash - -# On-Policy Distillation with Megatron-based teacher model -# This example uses the original model as the teacher (self-distillation for demonstration) -# -# IMPORTANT: This is just an example configuration! -# In practice, you should: -# 1. Use a different (stronger) model as the teacher -# 2. Adjust --opd-kl-coef based on your task -# 3. Configure proper evaluation metrics - -set -ex - -export PYTHONBUFFERED=16 - -NVLINK_COUNT=$(nvidia-smi topo -m 2>/dev/null | grep -o 'NV[0-9][0-9]*' | wc -l) -if [ "$NVLINK_COUNT" -gt 0 ]; then - HAS_NVLINK=1 -else - HAS_NVLINK=0 -fi -echo "HAS_NVLINK: $HAS_NVLINK (detected $NVLINK_COUNT NVLink references)" - -source "/root/vime/scripts/models/qwen3-8B.sh" - - -CKPT_ARGS=( - --hf-checkpoint /root/Qwen3-8B - --ref-load /root/Qwen3-8B_torch_dist - --load /root/Qwen3-8B_slime/ - --save /root/Qwen3-8B_slime/ - --save-interval 20 -) - -ROLLOUT_ARGS=( - --prompt-data /root/dapo-math-17k/dapo-math-17k.jsonl - --input-key prompt - --apply-chat-template - --rollout-shuffle - --num-rollout 300 - --rollout-batch-size 16 - --n-samples-per-prompt 4 - --rollout-max-response-len 16384 - --rollout-temperature 1 - - --global-batch-size 64 - --balance-data -) - -RM_ARGS=( - --rm-type math -) - -EVAL_ARGS=( - # --eval-interval 20 - # --eval-prompt-data aime ${DATA_DIR}/aime-2024/aime-2024.jsonl - # --n-samples-per-eval-prompt 16 - # --eval-max-response-len 16384 - # --eval-top-p 1 -) - -PERF_ARGS=( - --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 - - # --micro-batch-size 1 - --use-dynamic-batch-size - --max-tokens-per-gpu 16384 -) - -GRPO_ARGS=( - --advantage-estimator grpo # Base advantage estimator (can be ppo, grpo, etc.) - - # OPD Configuration - --use-opd # Enable on-policy distillation - --opd-type megatron # Use Megatron forward for teacher - --opd-kl-coef 1.0 # CHANGE THIS: KL penalty coefficient - # Teacher model configuration (CHANGE THIS to a stronger model!) - --opd-teacher-load /root/Qwen3-8B_torch_dist # Teacher model path - - --use-kl-loss - --kl-loss-coef 0.00 - --kl-loss-type low_var_kl - --entropy-coef 0.00 -) - -OPTIMIZER_ARGS=( - --optimizer adam - --lr 1e-6 - --lr-decay-style constant - --weight-decay 0.1 - --adam-beta1 0.9 - --adam-beta2 0.98 -) - -WANDB_ARGS=( - #--use-wandb - # --wandb-project vime-dev - # --wandb-group qwen3-8B-opd-megatron - # --wandb-key ${WANDB_KEY} -) - -VLLM_ARGS=( - --rollout-num-gpus-per-engine 1 - --vllm-gpu-memory-utilization 0.4 -) - - -MISC_ARGS=( - --attention-dropout 0.0 - --hidden-dropout 0.0 - --accumulate-allreduce-grads-in-fp32 - --attention-softmax-in-fp32 - --attention-backend flash -) - - - - -# launch the master node of ray in container -export MASTER_ADDR=${MASTER_ADDR:-"127.0.0.1"} -ray start --head --node-ip-address ${MASTER_ADDR} --num-gpus 8 --disable-usage-stats --dashboard-host=0.0.0.0 --dashboard-port=8265 - - -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 train.py \ - --actor-num-nodes 1 \ - --actor-num-gpus-per-node 2 \ - --rollout-num-gpus 4 \ - ${MODEL_ARGS[@]} \ - ${CKPT_ARGS[@]} \ - ${ROLLOUT_ARGS[@]} \ - ${OPTIMIZER_ARGS[@]} \ - ${GRPO_ARGS[@]} \ - ${WANDB_ARGS[@]} \ - ${PERF_ARGS[@]} \ - ${EVAL_ARGS[@]} \ - ${VLLM_ARGS[@]} \ - ${MISC_ARGS[@]} \ - ${RM_ARGS[@]} - - - -####clear after training -pkill -9 -f "vllm serve" -sleep 3 -ray stop --force -pkill -9 ray -pkill -9 python -sleep 3 -pkill -9 ray -pkill -9 python \ No newline at end of file diff --git a/examples/on_policy_distillation/run-qwen3-8B-opd.sh b/examples/on_policy_distillation/run-qwen3-8B-opd.sh deleted file mode 100644 index 2181235f1..000000000 --- a/examples/on_policy_distillation/run-qwen3-8B-opd.sh +++ /dev/null @@ -1,190 +0,0 @@ -#!/bin/bash - -# usage: bash examples/on_policy_distillation/run-qwen3-8B-opd.sh - -set -ex - - -# Start the teacher model server -TEACHER_IP="127.0.0.1" # Use localhost here, you can change it to your IP -TEACHER_PORT=13141 -LOG_FILE="/tmp/vllm_teacher_$(head /dev/urandom | tr -dc A-Za-z0-9 | head -c 6).log" - -## Launch the teacher model server in the background. -## OPD teacher uses /v1/completions with echo=True + prompt_logprobs=1 set -## per-request (see vime/rollout/on_policy_distillation.py reward_func). -## prompt_logprobs is a per-request SamplingParams field in vLLM; no -## server-side flag gates it. --max-logprobs defaults to 20, which is -## plenty for prompt_logprobs=1. -CUDA_VISIBLE_DEVICES=7 vllm serve /root/Qwen3-32B \ - --host 0.0.0.0 \ - --port $TEACHER_PORT \ - --tensor-parallel-size 1 \ - --gpu-memory-utilization 0.6 \ - --max-num-batched-tokens 4096 \ - > "$LOG_FILE" 2>&1 & - -echo "Starting teacher model server..." - -## Wait for the teacher model server to be ready -until curl -sf http://$TEACHER_IP:$TEACHER_PORT/health > /dev/null; do - echo "Waiting for the teacher model server to start..." - tail -n 10 "$LOG_FILE" - sleep 5 -done - -curl http://$TEACHER_IP:$TEACHER_PORT/v1/models -echo "Teacher model server is up and running at $TEACHER_IP:$TEACHER_PORT." -sleep 10 - - -export PYTHONBUFFERED=16 - -NVLINK_COUNT=$(nvidia-smi topo -m 2>/dev/null | grep -o 'NV[0-9][0-9]*' | wc -l) -if [ "$NVLINK_COUNT" -gt 0 ]; then - HAS_NVLINK=1 -else - HAS_NVLINK=0 -fi -echo "HAS_NVLINK: $HAS_NVLINK (detected $NVLINK_COUNT NVLink references)" - -source "/root/vime/scripts/models/qwen3-8B.sh" - - -CKPT_ARGS=( - --hf-checkpoint /root/Qwen3-8B - --ref-load /root/Qwen3-8B_torch_dist - --load /root/Qwen3-8B_slime/ - --save /root/Qwen3-8B_slime/ - --save-interval 20 -) - -ROLLOUT_ARGS=( - --prompt-data /root/dapo-math-17k/dapo-math-17k.jsonl - --input-key prompt - --apply-chat-template - --rollout-shuffle - --num-rollout 300 - --rollout-batch-size 16 - --n-samples-per-prompt 4 - --rollout-max-response-len 16384 - --rollout-temperature 1 - - --global-batch-size 64 - --balance-data -) - -RM_ARGS=( - --custom-rm-path vime.rollout.on_policy_distillation.reward_func - --custom-reward-post-process-path vime.rollout.on_policy_distillation.post_process_rewards - --rm-url http://$TEACHER_IP:$TEACHER_PORT/v1/completions -) - -EVAL_ARGS=( - # --eval-interval 20 - # --eval-prompt-data aime ${DATA_DIR}/aime-2024/aime-2024.jsonl - # --n-samples-per-eval-prompt 16 - # --eval-max-response-len 16384 - # --eval-top-p 1 -) - -PERF_ARGS=( - --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 - - # --micro-batch-size 1 - --use-dynamic-batch-size - --max-tokens-per-gpu 16384 -) - -GRPO_ARGS=( - --advantage-estimator grpo - --use-opd - --opd-type vllm - --opd-kl-coef 1.0 - --use-kl-loss - --kl-loss-coef 0.00 - --kl-loss-type low_var_kl - --entropy-coef 0.00 -) - -OPTIMIZER_ARGS=( - --optimizer adam - --lr 1e-6 - --lr-decay-style constant - --weight-decay 0.1 - --adam-beta1 0.9 - --adam-beta2 0.98 -) - -WANDB_ARGS=( - #--use-wandb - # --wandb-project vime-dev - # --wandb-group qwen3-8B-test - # --wandb-key ${WANDB_KEY} -) - -VLLM_ARGS=( - --rollout-num-gpus-per-engine 1 - --vllm-gpu-memory-utilization 0.4 -) - - -MISC_ARGS=( - --attention-dropout 0.0 - --hidden-dropout 0.0 - --accumulate-allreduce-grads-in-fp32 - --attention-softmax-in-fp32 - --attention-backend flash -) - - - - -# launch the master node of ray in container -export MASTER_ADDR=${MASTER_ADDR:-"127.0.0.1"} -ray start --head --node-ip-address ${MASTER_ADDR} --num-gpus 8 --disable-usage-stats --dashboard-host=0.0.0.0 --dashboard-port=8265 - - -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 train.py \ - --actor-num-nodes 1 \ - --actor-num-gpus-per-node 2 \ - --rollout-num-gpus 4 \ - ${MODEL_ARGS[@]} \ - ${CKPT_ARGS[@]} \ - ${ROLLOUT_ARGS[@]} \ - ${OPTIMIZER_ARGS[@]} \ - ${GRPO_ARGS[@]} \ - ${WANDB_ARGS[@]} \ - ${PERF_ARGS[@]} \ - ${EVAL_ARGS[@]} \ - ${VLLM_ARGS[@]} \ - ${MISC_ARGS[@]} \ - ${RM_ARGS[@]} - - - -####clear after training -pkill -9 -f "vllm serve" -sleep 3 -ray stop --force -pkill -9 ray -pkill -9 python -sleep 3 -pkill -9 ray -pkill -9 python \ No newline at end of file diff --git a/scripts/low_precision/run-kimi-k2-Thinking-int4.sh b/scripts/low_precision/run-kimi-k2-Thinking-int4.sh deleted file mode 100644 index 26ba9692e..000000000 --- a/scripts/low_precision/run-kimi-k2-Thinking-int4.sh +++ /dev/null @@ -1,182 +0,0 @@ -#!/bin/bash - -# for rerun the task -pkill -9 -f "vllm serve" -sleep 3 -ray stop --force -pkill -9 ray -pkill -9 python -sleep 3 -pkill -9 ray -pkill -9 python - -set -ex - -# will prevent ray from buffering stdout/stderr -export PYTHONBUFFERED=16 - -NVLINK_COUNT=$(nvidia-smi | grep -o "NVLink" | wc -l) -if [ "$NVLINK_COUNT" -gt 0 ]; then - HAS_NVLINK=1 -else - HAS_NVLINK=0 -fi -echo "HAS_NVLINK: $HAS_NVLINK (detected $NVLINK_COUNT NVLink references)" - -SCRIPT_DIR="$(cd -- "$(dirname -- "${BASH_SOURCE[0]}")" &>/dev/null && pwd)" -source "${SCRIPT_DIR}/../models/kimi-k2-thinking.sh" - -CKPT_ARGS=( - --hf-checkpoint /root/Kimi-K2-Thinking/ - --ref-load /root/Kimi-K2_thinking_torch_dist/ - --load /root/Kimi-K2-thinking_slime/ - --save /root/Kimi-K2-thinking_slime/ - --save-interval 20 -) - -ROLLOUT_ARGS=( - --prompt-data /root/dapo-math-17k/dapo-math-17k.jsonl - --input-key prompt - --label-key label - --apply-chat-template - --rollout-shuffle - - --rm-type math - - --num-rollout 100 - --rollout-batch-size 128 - --n-samples-per-prompt 8 - --rollout-max-response-len 16384 - --rollout-temperature 0.8 - - # --global-batch-size 256 - - --over-sampling-batch-size 256 - --dynamic-sampling-filter-path vime.rollout.filter_hub.dynamic_sampling_filters.check_reward_nonzero_std - - --num-steps-per-rollout 4 - --balance-data -) - -EVAL_ARGS=( - --eval-interval 10 - --eval-prompt-data aime /root/aime-2024/aime-2024.jsonl - --n-samples-per-eval-prompt 16 - --eval-max-response-len 16384 - --eval-top-p 0.7 -) - -PERF_ARGS=( - --tensor-model-parallel-size 8 - --sequence-parallel - --pipeline-model-parallel-size 8 - --context-parallel-size 4 - --expert-model-parallel-size 32 - --expert-tensor-parallel-size 1 - --decoder-last-pipeline-num-layers 5 - - --recompute-granularity full - --recompute-method uniform - --recompute-num-layers 1 - - --use-dynamic-batch-size - --max-tokens-per-gpu 16384 -) - -GRPO_ARGS=( - --advantage-estimator grpo - --use-kl-loss - --kl-loss-coef 0.00 - --kl-loss-type low_var_kl - # --kl-coef 0.00 - --entropy-coef 0.00 - --eps-clip 0.2 - --eps-clip-high 0.28 - --use-tis -) - -OPTIMIZER_ARGS=( - --optimizer adam - --lr 1e-6 - - --lr-decay-style constant - --weight-decay 0.1 - --adam-beta1 0.9 - --adam-beta2 0.98 - - --optimizer-cpu-offload - --overlap-cpu-optimizer-d2h-h2d - --use-precision-aware-optimizer -) - -WANDB_ARGS=( - # --use-wandb - # --wandb-project vime-dev - # --wandb-group kimi-k2-thinking-test - # --wandb-key ${WANDB_KEY} -) - -VLLM_ARGS=( - --rollout-num-gpus-per-engine 8 - --vllm-gpu-memory-utilization 0.7 - - # uncomment to enable DP-attention (non-MLA models only): - # --vllm-data-parallel-size 8 - - # uncomment to enable DeepEP all-to-all backend for MoE: - # --vllm-all2all-backend deepep_high_throughput - - # make every dp rank has 128 concurrency - --vllm-server-concurrency 1024 - --vllm-enable-expert-parallel -) - - -MISC_ARGS=( - # default dropout in megatron is 0.1 - --attention-dropout 0.0 - --hidden-dropout 0.0 - # should be good for model performance - --accumulate-allreduce-grads-in-fp32 - --attention-softmax-in-fp32 - # need to comment this when using model with MLA - --attention-backend flash - - # use deepep for megatron - # --moe-enable-deepep - # --moe-token-dispatcher-type flex - --no-check-for-nan-in-loss-and-grad -) - -# Build the runtime environment JSON with proper variable substitution -RUNTIME_ENV_JSON="{ - \"env_vars\": { - \"PYTHONPATH\": \"/root/Megatron-LM/\", - \"CUDA_DEVICE_MAX_CONNECTIONS\": \"1\", - \"NCCL_NVLS_ENABLE\": \"${HAS_NVLINK}\", - \"NCCL_TIMEOUT_MS\":\"360000000\", - \"no_proxy\": \"${no_proxy}\", - \"MASTER_ADDR\": \"${MASTER_ADDR}\", - \"OPEN_TRAINING_INT4_FAKE_QAT_FLAG\": \"1\", - \"OPEN_TRAINING_INT4_GROUP_SIZE\": \"32\" - } -}" - - -ray job submit --address="http://127.0.0.1:8265" \ - --runtime-env-json="${RUNTIME_ENV_JSON}" \ - -- python3 /personal/vime/vime/train.py \ - --actor-num-nodes 32 \ - --actor-num-gpus-per-node 8 \ - --colocate \ - --update-weight-buffer-size $(( 4 * 512 * 1024 * 1024)) \ - ${MODEL_ARGS[@]} \ - ${CKPT_ARGS[@]} \ - ${ROLLOUT_ARGS[@]} \ - ${OPTIMIZER_ARGS[@]} \ - ${GRPO_ARGS[@]} \ - ${WANDB_ARGS[@]} \ - ${PERF_ARGS[@]} \ - ${EVAL_ARGS[@]} \ - ${VLLM_ARGS[@]} \ - ${MISC_ARGS[@]} \ No newline at end of file diff --git a/scripts/low_precision/run-moonlight-16B-A3B-int4.sh b/scripts/low_precision/run-moonlight-16B-A3B-int4.sh deleted file mode 100644 index 655999103..000000000 --- a/scripts/low_precision/run-moonlight-16B-A3B-int4.sh +++ /dev/null @@ -1,165 +0,0 @@ -#!/bin/bash - -# for rerun the task -pkill -9 -f "vllm serve" -sleep 3 -ray stop --force -pkill -9 ray -pkill -9 python -sleep 3 -pkill -9 ray -pkill -9 python -pkill -9 redis - -set -ex - -# will prevent ray from buffering stdout/stderr -export PYTHONBUFFERED=16 - -NVLINK_COUNT=$(nvidia-smi topo -m 2>/dev/null | grep -o 'NV[0-9][0-9]*' | wc -l) -if [ "$NVLINK_COUNT" -gt 0 ]; then - HAS_NVLINK=1 -else - HAS_NVLINK=0 -fi -echo "HAS_NVLINK: $HAS_NVLINK (detected $NVLINK_COUNT NVLink references)" - -SCRIPT_DIR="$(cd -- "$(dirname -- "${BASH_SOURCE[0]}")" &>/dev/null && pwd)" -source "${SCRIPT_DIR}/../models/moonlight.sh" - -CKPT_ARGS=( - --hf-checkpoint /root/Moonlight-16B-A3B-Instruct-INT4 - --ref-load /root/Moonlight-16B-A3B-Instruct-INT4_torch_dist - --load /root/Moonlight-16B-A3B_slime/ - --save /root/Moonlight-16B-A3B_slime/ - --save-interval 20 -) - -ROLLOUT_ARGS=( - --prompt-data /root/dapo-math-17k/dapo-math-17k.jsonl - --input-key prompt - --label-key label - --apply-chat-template - --rollout-shuffle - --rm-type math - --num-rollout 3000 - --rollout-batch-size 128 - --n-samples-per-prompt 8 - --rollout-max-response-len 4096 - --rollout-temperature 0.8 - - --over-sampling-batch-size 256 - --dynamic-sampling-filter-path vime.rollout.filter_hub.dynamic_sampling_filters.check_reward_nonzero_std - - --num-steps-per-rollout 4 - # --global-batch-size 256 - --balance-data -) - -EVAL_ARGS=( - --eval-interval 20 - --eval-prompt-data aime /root/aime-2024/aime-2024.jsonl - --n-samples-per-eval-prompt 8 - --eval-max-response-len 4096 - --eval-top-p 0.7 -) - -PERF_ARGS=( - --tensor-model-parallel-size 2 - --sequence-parallel - --pipeline-model-parallel-size 1 - --context-parallel-size 1 - --expert-model-parallel-size 4 - --expert-tensor-parallel-size 1 - - --recompute-granularity full - --recompute-method uniform - --recompute-num-layers 1 - - # --micro-batch-size 1 - --use-dynamic-batch-size - --max-tokens-per-gpu 8192 -) - -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 -) - -OPTIMIZER_ARGS=( - --optimizer adam - --lr 1e-6 - --lr-decay-style constant - --weight-decay 0.1 - --adam-beta1 0.9 - --adam-beta2 0.98 - - --optimizer-cpu-offload - --overlap-cpu-optimizer-d2h-h2d - --use-precision-aware-optimizer -) - -WANDB_ARGS=( - # --use-wandb - # --wandb-project vime-dev - # --wandb-group moomlight-16B-A3B-test - # --wandb-key ${WANDB_KEY} -) - -VLLM_ARGS=( - --rollout-num-gpus-per-engine 4 - --vllm-gpu-memory-utilization 0.7 - --vllm-cudagraph-capture-sizes 1 2 4 8 $(seq 16 8 256) -) - -MISC_ARGS=( - # default dropout in megatron is 0.1 - --attention-dropout 0.0 - --hidden-dropout 0.0 - # should be good for model performance - --accumulate-allreduce-grads-in-fp32 - --attention-softmax-in-fp32 - # need to comment this when using model with MLA - # --attention-backend flash - - # use deepep for megatron - --moe-enable-deepep - --moe-token-dispatcher-type flex -) - -# launch the master node of ray in container -export MASTER_ADDR=${MASTER_ADDR:-"127.0.0.1"} -ray start --head --node-ip-address ${MASTER_ADDR} --num-gpus 4 --disable-usage-stats --dashboard-host=0.0.0.0 --dashboard-port=8265 - -# Build the runtime environment JSON with proper variable substitution -RUNTIME_ENV_JSON="{ - \"env_vars\": { - \"PYTHONPATH\": \"/root/Megatron-LM/\", - \"CUDA_DEVICE_MAX_CONNECTIONS\": \"1\", - \"NCCL_NVLS_ENABLE\": \"${HAS_NVLINK}\", - \"OPEN_TRAINING_INT4_FAKE_QAT_FLAG\": \"1\", - \"OPEN_TRAINING_INT4_GROUP_SIZE\": \"128\" - } -}" - -ray job submit --address="http://127.0.0.1:8265" \ - --runtime-env-json="${RUNTIME_ENV_JSON}" \ - -- python3 train.py \ - --actor-num-nodes 1 \ - --actor-num-gpus-per-node 4 \ - --colocate \ - ${MODEL_ARGS[@]} \ - ${CKPT_ARGS[@]} \ - ${ROLLOUT_ARGS[@]} \ - ${OPTIMIZER_ARGS[@]} \ - ${GRPO_ARGS[@]} \ - ${WANDB_ARGS[@]} \ - ${PERF_ARGS[@]} \ - ${EVAL_ARGS[@]} \ - ${VLLM_ARGS[@]} \ - ${MISC_ARGS[@]} diff --git a/scripts/low_precision/run-qwen3-235B-A22B-int4.sh b/scripts/low_precision/run-qwen3-235B-A22B-int4.sh deleted file mode 100644 index e278c3b75..000000000 --- a/scripts/low_precision/run-qwen3-235B-A22B-int4.sh +++ /dev/null @@ -1,169 +0,0 @@ -#!/bin/bash - -# for rerun the task -pkill -9 -f "vllm serve" -sleep 3 -ray stop --force -pkill -9 ray -pkill -9 python -sleep 3 -pkill -9 ray -pkill -9 python - -set -ex - -# will prevent ray from buffering stdout/stderr -export PYTHONBUFFERED=16 - -NVLINK_COUNT=$(nvidia-smi | grep -o "NVLink" | wc -l) -if [ "$NVLINK_COUNT" -gt 0 ]; then - HAS_NVLINK=1 -else - HAS_NVLINK=0 -fi -echo "HAS_NVLINK: $HAS_NVLINK (detected $NVLINK_COUNT NVLink references)" - -SCRIPT_DIR="$(cd -- "$(dirname -- "${BASH_SOURCE[0]}")" &>/dev/null && pwd)" -source "${SCRIPT_DIR}/../models/qwen3-235B-A22B.sh" - -CKPT_ARGS=( - --hf-checkpoint /root/Qwen3-235B-A22B-INT4/ - --ref-load /root/Qwen3-235B-A22B_torch_dist/ - --load /root/Qwen3-235B-A22B-vime/ - --save /root/Qwen3-235B-A22B-vime/ - --save-interval 20 -) - -ROLLOUT_ARGS=( - --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 300 - --rollout-batch-size 32 - --n-samples-per-prompt 8 - --rollout-max-response-len 8192 - --rollout-temperature 0.8 - - --global-batch-size 256 - --balance-data -) - -EVAL_ARGS=( - --eval-interval 10 - --eval-prompt-data aime /root/aime-2024/aime-2024.jsonl - --n-samples-per-eval-prompt 16 - --eval-max-response-len 16384 - --eval-top-p 0.7 -) - -PERF_ARGS=( - --tensor-model-parallel-size 4 - --sequence-parallel - --pipeline-model-parallel-size 4 - --context-parallel-size 2 - --expert-model-parallel-size 16 - --expert-tensor-parallel-size 1 - --decoder-last-pipeline-num-layers 22 - - --recompute-granularity full - --recompute-method uniform - --recompute-num-layers 1 - - # --micro-batch-size 1 - --use-dynamic-batch-size - --max-tokens-per-gpu 16384 -) - -GRPO_ARGS=( - --advantage-estimator grpo - --use-kl-loss - --kl-loss-coef 0.00 - --kl-loss-type low_var_kl - # --kl-coef 0.00 - --entropy-coef 0.00 - --eps-clip 0.2 - --eps-clip-high 0.28 - --use-tis -) - -OPTIMIZER_ARGS=( - --optimizer adam - --lr 1e-6 - --lr-decay-style constant - --weight-decay 0.1 - --adam-beta1 0.9 - --adam-beta2 0.98 - - --optimizer-cpu-offload - --overlap-cpu-optimizer-d2h-h2d - --use-precision-aware-optimizer -) - -WANDB_ARGS=( - # --use-wandb - # --wandb-project vime-dev - # --wandb-group qwen3-235B-A22B-test - # --wandb-key ${WANDB_KEY} -) - - -VLLM_ARGS=( - --rollout-num-gpus-per-engine 8 - --vllm-gpu-memory-utilization 0.7 - # # --vllm-data-parallel-size 4 - --vllm-enable-expert-parallel - --vllm-cudagraph-capture-sizes 1 2 4 8 $(seq 16 8 256) -) - - -MISC_ARGS=( - # default dropout in megatron is 0.1 - --attention-dropout 0.0 - --hidden-dropout 0.0 - # should be good for model performance - --accumulate-allreduce-grads-in-fp32 - --attention-softmax-in-fp32 - # need to comment this when using model with MLA - --attention-backend flash - --no-check-for-nan-in-loss-and-grad - - # use deepep for megatron - # --moe-enable-deepep - # --moe-token-dispatcher-type flex -) - -# Build the runtime environment JSON with proper variable substitution -RUNTIME_ENV_JSON="{ - \"env_vars\": { - \"PYTHONPATH\": \"/root/Megatron-LM/\", - \"CUDA_DEVICE_MAX_CONNECTIONS\": \"1\", - \"NCCL_NVLS_ENABLE\": \"${HAS_NVLINK}\", - \"NCCL_TIMEOUT_MS\":\"360000000\", - \"no_proxy\": \"${no_proxy}\", - \"MASTER_ADDR\": \"${MASTER_ADDR}\", - \"OPEN_TRAINING_INT4_FAKE_QAT_FLAG\": \"1\", - \"OPEN_TRAINING_INT4_GROUP_SIZE\": \"128\" - } -}" - -ray job submit --address="http://127.0.0.1:8265" \ - --runtime-env-json="${RUNTIME_ENV_JSON}" \ - -- python3 train.py \ - --actor-num-nodes 8 \ - --actor-num-gpus-per-node 8 \ - --colocate \ - ${MODEL_ARGS[@]} \ - ${CKPT_ARGS[@]} \ - ${ROLLOUT_ARGS[@]} \ - ${OPTIMIZER_ARGS[@]} \ - ${GRPO_ARGS[@]} \ - ${WANDB_ARGS[@]} \ - ${PERF_ARGS[@]} \ - ${EVAL_ARGS[@]} \ - ${VLLM_ARGS[@]} \ - ${MISC_ARGS[@]} \ No newline at end of file diff --git a/scripts/low_precision/run-qwen3-30B-A3B-int4.sh b/scripts/low_precision/run-qwen3-30B-A3B-int4.sh deleted file mode 100644 index c927cd59e..000000000 --- a/scripts/low_precision/run-qwen3-30B-A3B-int4.sh +++ /dev/null @@ -1,164 +0,0 @@ -#!/bin/bash - -# for rerun the task -pkill -9 -f "vllm serve" -sleep 3 -ray stop --force -pkill -9 ray -pkill -9 python -sleep 3 -pkill -9 ray -pkill -9 python - -set -ex - -# will prevent ray from buffering stdout/stderrs -export PYTHONBUFFERED=16 - -NVLINK_COUNT=$(nvidia-smi topo -m 2>/dev/null | grep -o 'NV[0-9][0-9]*' | wc -l) -if [ "$NVLINK_COUNT" -gt 0 ]; then - HAS_NVLINK=1 -else - HAS_NVLINK=0 -fi -echo "HAS_NVLINK: $HAS_NVLINK (detected $NVLINK_COUNT NVLink references)" - -SCRIPT_DIR="$(cd -- "$(dirname -- "${BASH_SOURCE[0]}")" &>/dev/null && pwd)" -source "${SCRIPT_DIR}/../models/qwen3-30B-A3B.sh" - -CKPT_ARGS=( - --hf-checkpoint /root/Qwen3-30B-A3B-INT4/ - --ref-load /root/Qwen3-30B-A3B_torch_dist/ - --load /root/Qwen3-30B-A3B_slime/ - --save /root/Qwen3-30B-A3B_slime/ - --save-interval 20 -) - -ROLLOUT_ARGS=( - --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 100 - --rollout-batch-size 32 - --n-samples-per-prompt 8 - --rollout-max-response-len 8192 - --rollout-temperature 0.8 - - --global-batch-size 256 - --balance-data - # --debug-rollout-only -) - -EVAL_ARGS=( - --eval-interval 10 - --eval-prompt-data /root/aime-2024/aime-2024.jsonl - --n-samples-per-eval-prompt 8 - --eval-max-response-len 16384 - --eval-top-p 0.7 -) - -PERF_ARGS=( - --tensor-model-parallel-size 4 - --sequence-parallel - --pipeline-model-parallel-size 1 - --context-parallel-size 1 - --expert-model-parallel-size 8 - --expert-tensor-parallel-size 1 - - --recompute-granularity full - --recompute-method uniform - --recompute-num-layers 1 - - # --micro-batch-size 1 - --use-dynamic-batch-size - --max-tokens-per-gpu 8192 -) - -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 - --use-tis -) - -OPTIMIZER_ARGS=( - --optimizer adam - --lr 1e-6 - --lr-decay-style constant - --weight-decay 0.1 - --adam-beta1 0.9 - --adam-beta2 0.98 - - --optimizer-cpu-offload - --overlap-cpu-optimizer-d2h-h2d - --use-precision-aware-optimizer -) - -WANDB_ARGS=( - # --use-wandb - # --wandb-project vime-dev - # --wandb-group qwen3-30B-A3B-test - # --wandb-key ${WANDB_KEY} -) - -VLLM_ARGS=( - --rollout-num-gpus-per-engine 1 - --vllm-gpu-memory-utilization 0.7 - --vllm-cudagraph-capture-sizes 1 2 4 8 $(seq 16 8 256) -) - -MISC_ARGS=( - # default dropout in megatron is 0.1 - --attention-dropout 0.0 - --hidden-dropout 0.0 - # should be good for model performance - --accumulate-allreduce-grads-in-fp32 - --attention-softmax-in-fp32 - # need to comment this when using model with MLA - --attention-backend flash - # use deepep for megatron - # --moe-enable-deepep - # --moe-token-dispatcher-type flex -) - -# launch the master node of ray in container -export MASTER_ADDR=${MASTER_ADDR:-"127.0.0.1"} -ray start --head --node-ip-address ${MASTER_ADDR} --num-gpus 8 --disable-usage-stats --dashboard-host=0.0.0.0 --dashboard-port=8265 - -# Build the runtime environment JSON with proper variable substitution -RUNTIME_ENV_JSON="{ - \"env_vars\": { - \"PYTHONPATH\": \"/root/Megatron-LM/\", - \"CUDA_DEVICE_MAX_CONNECTIONS\": \"1\", - \"NCCL_NVLS_ENABLE\": \"${HAS_NVLINK}\", - \"OPEN_TRAINING_INT4_FAKE_QAT_FLAG\": \"1\", - \"OPEN_TRAINING_INT4_GROUP_SIZE\": \"128\" - } -}" - -ray job submit --address="http://127.0.0.1:8265" \ - --runtime-env-json="${RUNTIME_ENV_JSON}" \ - -- python3 train.py \ - --actor-num-nodes 1 \ - --actor-num-gpus-per-node 8 \ - --colocate \ - ${MODEL_ARGS[@]} \ - ${CKPT_ARGS[@]} \ - ${ROLLOUT_ARGS[@]} \ - ${OPTIMIZER_ARGS[@]} \ - ${GRPO_ARGS[@]} \ - ${WANDB_ARGS[@]} \ - ${PERF_ARGS[@]} \ - ${EVAL_ARGS[@]} \ - ${VLLM_ARGS[@]} \ - ${MISC_ARGS[@]} - \ No newline at end of file diff --git a/scripts/low_precision/run-qwen3-30b-a3b-fp8.sh b/scripts/low_precision/run-qwen3-30b-a3b-fp8.sh deleted file mode 100644 index 6f58057e7..000000000 --- a/scripts/low_precision/run-qwen3-30b-a3b-fp8.sh +++ /dev/null @@ -1,179 +0,0 @@ -#!/bin/bash - -# for rerun the task -# pkill -9 -f "vllm serve" -# sleep 3 -# ray stop --force -# pkill -9 ray -# pkill -9 python -# sleep 3 -# pkill -9 ray -# pkill -9 python -# pkill -9 redis - -set -ex - -# will prevent ray from buffering stdout/stderr -export PYTHONBUFFERED=16 - -NVLINK_COUNT=$(nvidia-smi topo -m 2>/dev/null | grep -o 'NV[0-9][0-9]*' | wc -l) -if [ "$NVLINK_COUNT" -gt 0 ]; then - HAS_NVLINK=1 -else - HAS_NVLINK=0 -fi -echo "HAS_NVLINK: $HAS_NVLINK (detected $NVLINK_COUNT NVLink references)" - -SCRIPT_DIR="$(cd -- "$(dirname -- "${BASH_SOURCE[0]}")" &>/dev/null && pwd)" -source "${SCRIPT_DIR}/../scripts/models/qwen3-30B-A3B.sh" - -# Base directory for checkpoints and related files (adjust if necessary) -BASE_DIR="/root" - -CKPT_ARGS=( - --hf-checkpoint "${BASE_DIR}/Qwen3-30B-A3B-FP8/" - --ref-load "${BASE_DIR}/Qwen3-30B-A3B_torch_dist/" - --load "${BASE_DIR}/Qwen3-30B-A3B_slime/" - --save "${BASE_DIR}/Qwen3-30B-A3B_slime/" - --save-interval 20 -) - -ROLLOUT_ARGS=( - --prompt-data "${BASE_DIR}/dapo-math-17k.jsonl" - --input-key prompt - --label-key label - --apply-chat-template - --rollout-shuffle - --rm-type math - --num-rollout 200 - --rollout-batch-size 16 - --n-samples-per-prompt 8 - --rollout-max-response-len 8192 - --rollout-temperature 1 - - --global-batch-size 128 - --balance-data -) - -EVAL_ARGS=( - --eval-interval 20 - --eval-prompt-data aime "${BASE_DIR}/aime-2024.jsonl" - --n-samples-per-eval-prompt 16 - --eval-max-response-len 16384 - --eval-top-p 1 -) - -PERF_ARGS=( - --tensor-model-parallel-size 1 - --sequence-parallel - --pipeline-model-parallel-size 4 - --context-parallel-size 1 - --expert-model-parallel-size 4 - --expert-tensor-parallel-size 1 - - --recompute-granularity full - --recompute-method uniform - --recompute-num-layers 1 - - # --micro-batch-size 1 - --use-dynamic-batch-size - --max-tokens-per-gpu 20480 - - # use deepep for megatron - --moe-enable-deepep - --moe-token-dispatcher-type flex - - # fp8 - --transformer-impl transformer_engine - --bf16 - --fp8-format e4m3 - --fp8-recipe blockwise - # --fp8-param-gather -) - -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 - --use-tis -) - -OPTIMIZER_ARGS=( - --optimizer adam - --lr 1e-6 - --lr-decay-style constant - --weight-decay 0.1 - --adam-beta1 0.9 - --adam-beta2 0.98 - - --optimizer-cpu-offload - --overlap-cpu-optimizer-d2h-h2d - --use-precision-aware-optimizer -) - -WANDB_ARGS=( - #--use-wandb - # --wandb-project vime-dev - # --wandb-group qwen3-30B-A3B-test - # --wandb-key ${WANDB_KEY} -) - -VLLM_ARGS=( - --rollout-num-gpus-per-engine 8 - --vllm-gpu-memory-utilization 0.6 - --vllm-enable-expert-parallel - --vllm-cudagraph-capture-sizes 1 2 4 8 $(seq 16 8 256) - # --use-rollout-routing-replay -) - -MISC_ARGS=( - # default dropout in megatron is 0.1 - --attention-dropout 0.0 - --hidden-dropout 0.0 - # should be good for model performance - --accumulate-allreduce-grads-in-fp32 - --attention-softmax-in-fp32 - # need to comment this when using model with MLA - --attention-backend flash -) - -# Get Ray Head node info automatically -ip=$(ps aux | grep dashboard | grep -oP '(?<=--node-ip-address=)[0-9\.]+' | head -1) -port=$(ps aux | grep dashboard | grep -oP '(?<=dashboard-port=)\d+' | head -1) -export HEAD_NODE_ADDRESS="$ip" -export DASHBOARD_PORT="$port" -echo "Detected Ray Head IP: $HEAD_NODE_ADDRESS, Port: $DASHBOARD_PORT" - -export RAY_ADDRESS="http://${HEAD_NODE_ADDRESS}:${DASHBOARD_PORT}" - -# You should enable NVTE_FP8_BLOCK_SCALING_FP32_SCALES to use fp32 scales in fp8 training -RUNTIME_ENV_JSON="{ - \"env_vars\": { - \"PYTHONPATH\": \"/root/Megatron-LM/\", - \"CUDA_DEVICE_MAX_CONNECTIONS\": \"1\", - \"NCCL_NVLS_ENABLE\": \"${HAS_NVLINK}\", - \"NVTE_FP8_BLOCK_SCALING_FP32_SCALES\": \"1\", - \"NCCL_TIMEOUT_MS\":\"36000000\" - } -}" - -ray job submit --address="${RAY_ADDRESS}" \ - --runtime-env-json="${RUNTIME_ENV_JSON}" \ - -- python3 train.py \ - --actor-num-nodes 2 \ - --actor-num-gpus-per-node 8 \ - --colocate \ - ${MODEL_ARGS[@]} \ - ${CKPT_ARGS[@]} \ - ${ROLLOUT_ARGS[@]} \ - ${OPTIMIZER_ARGS[@]} \ - ${GRPO_ARGS[@]} \ - ${WANDB_ARGS[@]} \ - ${PERF_ARGS[@]} \ - ${EVAL_ARGS[@]} \ - ${VLLM_ARGS[@]} \ - ${MISC_ARGS[@]} \ No newline at end of file diff --git a/scripts/low_precision/run-qwen3-4b-fp8.sh b/scripts/low_precision/run-qwen3-4b-fp8.sh deleted file mode 100644 index 3c24bfa22..000000000 --- a/scripts/low_precision/run-qwen3-4b-fp8.sh +++ /dev/null @@ -1,154 +0,0 @@ -#!/bin/bash - -# for rerun the task -pkill -9 -f "vllm serve" -sleep 3 -ray stop --force -pkill -9 ray -pkill -9 python -sleep 3 -pkill -9 ray -pkill -9 python - -set -ex - -# will prevent ray from buffering stdout/stderr -export PYTHONBUFFERED=16 - -NVLINK_COUNT=$(nvidia-smi topo -m 2>/dev/null | grep -o 'NV[0-9][0-9]*' | wc -l) -if [ "$NVLINK_COUNT" -gt 0 ]; then - HAS_NVLINK=1 -else - HAS_NVLINK=0 -fi -echo "HAS_NVLINK: $HAS_NVLINK (detected $NVLINK_COUNT NVLink references)" - -SCRIPT_DIR="$(cd -- "$(dirname -- "${BASH_SOURCE[0]}")" &>/dev/null && pwd)" -source "${SCRIPT_DIR}/../scripts/models/qwen3-4B.sh" - -CKPT_ARGS=( - --hf-checkpoint /root/Qwen3-4B-FP8 - #--hf-checkpoint /root/Qwen3-4B-FP8 - --ref-load /root/Qwen3-4B_torch_dist - --load /root/qwen3-4b_cp8_fp8 - --save /root/rl-model/qwen3-4b_cp8_fp8 - --save-interval 20 -) - -ROLLOUT_ARGS=( - --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_ARGS=( - --eval-interval 20 - --eval-prompt-data aime /root/data/aime-2024.jsonl - --n-samples-per-eval-prompt 16 - --eval-max-response-len 16384 - --eval-top-p 1 -) - -PERF_ARGS=( - --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 - - # --micro-batch-size 1 - --use-dynamic-batch-size - --max-tokens-per-gpu 9216 -) - -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 -) - -OPTIMIZER_ARGS=( - --optimizer adam - --lr 1e-6 - --lr-decay-style constant - --weight-decay 0.1 - --adam-beta1 0.9 - --adam-beta2 0.98 - -) - -WANDB_ARGS=( - # --use-wandb - # --wandb-project vime-dev - # --wandb-group qwen3-4B-test - # --wandb-key ${WANDB_KEY} -) - -VLLM_ARGS=( - --rollout-num-gpus-per-engine 2 - --vllm-gpu-memory-utilization 0.7 -) - -MISC_ARGS=( - # default dropout in megatron is 0.1 - --attention-dropout 0.0 - --hidden-dropout 0.0 - # should be good for model performance - --accumulate-allreduce-grads-in-fp32 - --attention-softmax-in-fp32 - # need to comment this when using model with MLA - --attention-backend flash -) - - -# launch the master node of ray in container -export MASTER_ADDR=${MASTER_ADDR:-"127.0.0.1"} -ray start --head --node-ip-address ${MASTER_ADDR} --num-gpus 8 --disable-usage-stats --dashboard-host=0.0.0.0 --dashboard-port=8265 - -# Build the runtime environment JSON with proper variable substitution -# you should enable NVTE_FP8_BLOCK_SCALING_FP32_SCALES to use fp32 scales in fp8 training -RUNTIME_ENV_JSON="{ - \"env_vars\": { - \"PYTHONPATH\": \"/root/Megatron-LM/:${SCRIPT_DIR}\", - \"CUDA_DEVICE_MAX_CONNECTIONS\": \"1\", - \"NCCL_NVLS_ENABLE\": \"${HAS_NVLINK}\", - \"NVTE_FP8_BLOCK_SCALING_FP32_SCALES\": \"1\" - } -}" - -ray job submit --address="http://127.0.0.1:8265" \ - --runtime-env-json="${RUNTIME_ENV_JSON}" \ - -- python3 train.py \ - --actor-num-nodes 1 \ - --actor-num-gpus-per-node 8 \ - --colocate \ - ${MODEL_ARGS[@]} \ - ${CKPT_ARGS[@]} \ - ${ROLLOUT_ARGS[@]} \ - ${OPTIMIZER_ARGS[@]} \ - ${GRPO_ARGS[@]} \ - ${WANDB_ARGS[@]} \ - ${PERF_ARGS[@]} \ - ${EVAL_ARGS[@]} \ - ${VLLM_ARGS[@]} \ - ${MISC_ARGS[@]} \ No newline at end of file diff --git a/scripts/models/mimo-7B-rl.sh b/scripts/models/mimo-7B-rl.sh index 22366935f..3def88224 100644 --- a/scripts/models/mimo-7B-rl.sh +++ b/scripts/models/mimo-7B-rl.sh @@ -15,5 +15,4 @@ MODEL_ARGS=( --vocab-size 151680 --untie-embeddings-and-output-weights --max-position-embeddings 32768 - --mtp-num-layers 1 ) diff --git a/scripts/run-deepseek-r1.sh b/scripts/run-deepseek-r1.sh deleted file mode 100644 index ea06d11f6..000000000 --- a/scripts/run-deepseek-r1.sh +++ /dev/null @@ -1,171 +0,0 @@ -#!/bin/bash - -# for rerun the task -pkill -9 -f "vllm serve" -sleep 3 -ray stop --force -pkill -9 ray -pkill -9 python -sleep 3 -pkill -9 ray -pkill -9 python - -set -ex - -# will prevent ray from buffering stdout/stderr -export PYTHONBUFFERED=16 - -NVLINK_COUNT=$(nvidia-smi topo -m 2>/dev/null | grep -o 'NV[0-9][0-9]*' | wc -l) -if [ "$NVLINK_COUNT" -gt 0 ]; then - HAS_NVLINK=1 -else - HAS_NVLINK=0 -fi -echo "HAS_NVLINK: $HAS_NVLINK (detected $NVLINK_COUNT NVLink references)" - -SCRIPT_DIR="$(cd -- "$(dirname -- "${BASH_SOURCE[0]}")" &>/dev/null && pwd)" -source "${SCRIPT_DIR}/models/deepseek-v3.sh" - -CKPT_ARGS=( - --hf-checkpoint $BASE_DIR/DeepSeek-R1/ - #--hf-checkpoint $BASE_DIR/DeepSeek-R1-bf16/ - --ref-load $BASE_DIR/DeepSeek-R1_torch_dist/ - --load $BASE_DIR/DeepSeek-R1_slime/ - --save $BASE_DIR/DeepSeek-R1_slime/ - --save-interval 20 -) - -ROLLOUT_ARGS=( - --prompt-data $BASE_DIR/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 128 - --n-samples-per-prompt 8 - --rollout-max-response-len 32768 - --rollout-temperature 1 - - --over-sampling-batch-size 256 - --dynamic-sampling-filter-path vime.rollout.filter_hub.dynamic_sampling_filters.check_reward_nonzero_std - - --num-steps-per-rollout 4 - --balance-data -) - -EVAL_ARGS=( - --eval-interval 20 - --eval-prompt-data aime $BASE_DIR/rl_data/aime-2024.jsonl - --n-samples-per-eval-prompt 8 - --eval-max-response-len 32768 - --eval-top-p 1 -) - -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=( - --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_ARGS=( - --optimizer adam - --lr 1e-6 - --lr-decay-style constant - --weight-decay 0.1 - --adam-beta1 0.9 - --adam-beta2 0.98 - - --optimizer-cpu-offload - --overlap-cpu-optimizer-d2h-h2d - --use-precision-aware-optimizer -) - -WANDB_ARGS=( - # --use-wandb - # --wandb-project vime-dev - # --wandb-group deepseek-r1-test - # --wandb-key ${WANDB_KEY} -) - -VLLM_ARGS=( - --rollout-num-gpus-per-engine 64 - --vllm-gpu-memory-utilization 0.7 - - # dp attention - --vllm-data-parallel-size 8 - - # mtp - - # make every dp rank has 128 concurrency - --vllm-server-concurrency 1024 - --vllm-enable-expert-parallel - --vllm-all2all-backend deepep_high_throughput - --vllm-speculative-config '{"method":"mtp","num_speculative_tokens":3}' -) - -MISC_ARGS=( - # default dropout in megatron is 0.1 - --attention-dropout 0.0 - --hidden-dropout 0.0 - # should be good for model performance - --accumulate-allreduce-grads-in-fp32 - --attention-softmax-in-fp32 - --attention-backend flash - - # use deepep for megatron - --moe-enable-deepep - --moe-token-dispatcher-type flex -) - -# launch the master node of ray in container -export no_proxy="127.0.0.1,${MASTER_ADDR}" -ray start --head --node-ip-address ${MASTER_ADDR} --num-gpus 8 --disable-usage-stats --dashboard-host=0.0.0.0 --dashboard-port=8265 - - -ray job submit --address="http://127.0.0.1:8265" \ - --runtime-env-json='{ - "env_vars": { - "no_proxy": "localhost,127.0.0.1,0.0.0.0,${MASTER_ADDR}", - "MASTER_ADDR": "${MASTER_ADDR}", - "PYTHONPATH": "/root/Megatron-LM/", - "CUDA_DEVICE_MAX_CONNECTIONS": "1", - "LD_LIBRARY_PATH": "/usr/local/nvidia/lib:/usr/local/nvidia/lib64" - } - }' \ - -- python3 train.py \ - --actor-num-nodes 16 \ - --actor-num-gpus-per-node 8 \ - --colocate \ - ${MODEL_ARGS[@]} \ - ${CKPT_ARGS[@]} \ - ${ROLLOUT_ARGS[@]} \ - ${OPTIMIZER_ARGS[@]} \ - ${GRPO_ARGS[@]} \ - ${WANDB_ARGS[@]} \ - ${PERF_ARGS[@]} \ - ${EVAL_ARGS[@]} \ - ${VLLM_ARGS[@]} \ - ${MISC_ARGS[@]} diff --git a/scripts/run-glm4-9B.sh b/scripts/run-glm4-9B.sh deleted file mode 100644 index ce69ac49c..000000000 --- a/scripts/run-glm4-9B.sh +++ /dev/null @@ -1,150 +0,0 @@ -#!/bin/bash - -# for rerun the task -pkill -9 -f "vllm serve" -sleep 3 -ray stop --force -pkill -9 ray -pkill -9 python -sleep 3 -pkill -9 ray -pkill -9 python - -set -ex - -# will prevent ray from buffering stdout/stderr -export PYTHONBUFFERED=16 - -NVLINK_COUNT=$(nvidia-smi topo -m 2>/dev/null | grep -o 'NV[0-9][0-9]*' | wc -l) -if [ "$NVLINK_COUNT" -gt 0 ]; then - HAS_NVLINK=1 -else - HAS_NVLINK=0 -fi -echo "HAS_NVLINK: $HAS_NVLINK (detected $NVLINK_COUNT NVLink references)" - -SCRIPT_DIR="$(cd -- "$(dirname -- "${BASH_SOURCE[0]}")" &>/dev/null && pwd)" -source "${SCRIPT_DIR}/models/glm4-9B.sh" - -CKPT_ARGS=( - --hf-checkpoint /root/GLM-Z1-9B-0414/ - --ref-load /root/GLM-Z1-9B-0414_torch_dist - --load /root/GLM-Z1-9B-0414_slime/ - --save /root/GLM-Z1-9B-0414_slime/ - --save-interval 20 -) - -ROLLOUT_ARGS=( - --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_ARGS=( - --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 -) - -PERF_ARGS=( - --tensor-model-parallel-size 2 - --sequence-parallel - --pipeline-model-parallel-size 1 - --context-parallel-size 2 - --expert-model-parallel-size 1 - --expert-tensor-parallel-size 1 - - --recompute-granularity full - --recompute-method uniform - --recompute-num-layers 1 - - # --micro-batch-size 1 - --use-dynamic-batch-size - --max-tokens-per-gpu 4608 -) - -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 -) - -OPTIMIZER_ARGS=( - --optimizer adam - --lr 1e-6 - --lr-decay-style constant - --weight-decay 0.1 - --adam-beta1 0.9 - --adam-beta2 0.98 -) - -WANDB_ARGS=( - #--use-wandb - # --wandb-project vime-dev - # --wandb-group qwen3-4B-test - # --wandb-key ${WANDB_KEY} -) - -VLLM_ARGS=( - --rollout-num-gpus-per-engine 2 -) - -MISC_ARGS=( - # default dropout in megatron is 0.1 - --attention-dropout 0.0 - --hidden-dropout 0.0 - # should be good for model performance - --accumulate-allreduce-grads-in-fp32 - --attention-softmax-in-fp32 - # need to comment this when using model with MLA - --attention-backend flash -) - -# launch the master node of ray in container -export MASTER_ADDR=${MASTER_ADDR:-"127.0.0.1"} -ray start --head --node-ip-address ${MASTER_ADDR} --num-gpus 8 --disable-usage-stats --dashboard-host=0.0.0.0 --dashboard-port=8265 - -# Build the runtime environment JSON with proper variable substitution -RUNTIME_ENV_JSON="{ - \"env_vars\": { - \"PYTHONPATH\": \"/root/Megatron-LM/\", - \"CUDA_DEVICE_MAX_CONNECTIONS\": \"1\", - \"NCCL_NVLS_ENABLE\": \"${HAS_NVLINK}\" - } -}" - -ray job submit --address="http://127.0.0.1:8265" \ - --runtime-env-json="${RUNTIME_ENV_JSON}" \ - -- python3 train.py \ - --actor-num-nodes 1 \ - --actor-num-gpus-per-node 4 \ - --rollout-num-gpus 4 \ - ${MODEL_ARGS[@]} \ - ${CKPT_ARGS[@]} \ - ${ROLLOUT_ARGS[@]} \ - ${OPTIMIZER_ARGS[@]} \ - ${GRPO_ARGS[@]} \ - ${WANDB_ARGS[@]} \ - ${PERF_ARGS[@]} \ - ${EVAL_ARGS[@]} \ - ${VLLM_ARGS[@]} \ - ${MISC_ARGS[@]} diff --git a/scripts/run-glm4.7-30B-A3B.sh b/scripts/run-glm4.7-30B-A3B.sh deleted file mode 100644 index 9cbcb947c..000000000 --- a/scripts/run-glm4.7-30B-A3B.sh +++ /dev/null @@ -1,171 +0,0 @@ -#!/bin/bash - -# for rerun the task -pkill -9 -f "vllm serve" -sleep 3 -ray stop --force -pkill -9 ray -pkill -9 python -sleep 3 -pkill -9 ray -pkill -9 python - -set -ex - -# will prevent ray from buffering stdout/stderr -export PYTHONBUFFERED=16 - -unset http_proxy https_proxy HTTP_PROXY HTTPS_PROXY - -NVLINK_COUNT=$(nvidia-smi topo -m 2>/dev/null | grep -o 'NV[0-9][0-9]*' | wc -l) -if [ "$NVLINK_COUNT" -gt 0 ]; then - HAS_NVLINK=1 -else - HAS_NVLINK=0 -fi -echo "HAS_NVLINK: $HAS_NVLINK (detected $NVLINK_COUNT NVLink references)" - -SCRIPT_DIR="$(cd -- "$(dirname -- "${BASH_SOURCE[0]}")" &>/dev/null && pwd)" -source "${SCRIPT_DIR}/../scripts/models/glm4.7-30B-A3B.sh" - -CKPT_ARGS=( - --hf-checkpoint $BASE_DIR/GLM-4.7-Flash - --ref-load $BASE_DIR/GLM-4.7-Flash_torch_dist/ -) - -ROLLOUT_ARGS=( - --prompt-data $BASE_DIR/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 128 - --n-samples-per-prompt 8 - --rollout-max-response-len 32768 - --rollout-temperature 1.0 - - --global-batch-size 1024 - #--balance-data -) - -EVAL_ARGS=( - --eval-interval 20 - --eval-prompt-data aime24 $BASE_DIR/rl_data/aime-2024.jsonl - --n-samples-per-eval-prompt 2 - --eval-max-response-len 32768 - --eval-temperature 1.0 - --eval-top-p 0.95 -) - -PERF_ARGS=( - --tensor-model-parallel-size 2 - --sequence-parallel - --pipeline-model-parallel-size 2 - --context-parallel-size 2 - --expert-model-parallel-size 8 - --expert-tensor-parallel-size 1 - --decoder-last-pipeline-num-layers 23 - - --recompute-granularity full - --recompute-method uniform - --recompute-num-layers 1 - - --use-dynamic-batch-size - --max-tokens-per-gpu 32768 -) - -MTP_ARGS=( - --mtp-num-layers 1 - --enable-mtp-training - --mtp-loss-scaling-factor 0.2 -) - -GRPO_ARGS=( - --advantage-estimator grpo - --use-kl-loss - --kl-loss-coef 0.00 - --kl-loss-type low_var_kl - --kl-coef 0.00 - --entropy-coef 0.00 -) - -OPTIMIZER_ARGS=( - --optimizer adam - --lr 1e-6 - --lr-decay-style constant - --weight-decay 0.1 - --adam-beta1 0.9 - --adam-beta2 0.98 - - --optimizer-cpu-offload - --overlap-cpu-optimizer-d2h-h2d - --use-precision-aware-optimizer -) - -WANDB_ARGS=( - # --use-wandb - # --wandb-project vime-dev - # --wandb-group glm4.7-flash -) - -VLLM_ARGS=( - --rollout-num-gpus-per-engine 8 - --vllm-gpu-memory-utilization 0.8 - --vllm-data-parallel-size 8 - - # mtp - - - --vllm-max-num-seqs 512 - --vllm-speculative-config '{"method":"mtp","num_speculative_tokens":3}' - --vllm-max-cudagraph-capture-size 64 -) - -MISC_ARGS=( - # default dropout in megatron is 0.1 - --attention-dropout 0.0 - --hidden-dropout 0.0 - # should be good for model performance - --accumulate-allreduce-grads-in-fp32 - --attention-softmax-in-fp32 - # need to comment this when using model with MLA - --attention-backend flash - - --moe-token-dispatcher-type flex - --moe-enable-deepep -) - -# launch the master node of ray in container -export MASTER_ADDR=${MASTER_ADDR:-"127.0.0.1"} -ray start --head --node-ip-address ${MASTER_ADDR} --num-gpus 8 --disable-usage-stats --dashboard-host=0.0.0.0 --dashboard-port=8265 - -# Build the runtime environment JSON with proper variable substitution -RUNTIME_ENV_JSON="{ - \"env_vars\": { - \"PYTHONPATH\": \"/root/Megatron-LM/\", - \"CUDA_DEVICE_MAX_CONNECTIONS\": \"1\", - \"NCCL_NVLS_ENABLE\": \"${HAS_NVLINK}\" - } -}" - -ray job submit --address="http://127.0.0.1:8265" \ - --runtime-env-json="${RUNTIME_ENV_JSON}" \ - -- python3 train.py \ - --actor-num-nodes 1 \ - --actor-num-gpus-per-node 8 \ - --colocate \ - ${MODEL_ARGS[@]} \ - ${CKPT_ARGS[@]} \ - ${ROLLOUT_ARGS[@]} \ - ${OPTIMIZER_ARGS[@]} \ - ${GRPO_ARGS[@]} \ - ${WANDB_ARGS[@]} \ - ${PERF_ARGS[@]} \ - ${EVAL_ARGS[@]} \ - ${VLLM_ARGS[@]} \ - ${MISC_ARGS[@]} \ - ${MTP_ARGS[@]} diff --git a/scripts/run-glm4.7-355B-A32B.sh b/scripts/run-glm4.7-355B-A32B.sh deleted file mode 100644 index c7c9ca9ae..000000000 --- a/scripts/run-glm4.7-355B-A32B.sh +++ /dev/null @@ -1,194 +0,0 @@ -#!/bin/bash - -# for rerun the task -pkill -9 -f "vllm serve" -sleep 3 -ray stop --force -pkill -9 ray -pkill -9 python -sleep 3 -pkill -9 ray -pkill -9 python - -set -ex - -# will prevent ray from buffering stdout/stderr -export PYTHONBUFFERED=16 - -unset http_proxy https_proxy HTTP_PROXY HTTPS_PROXY - -NVLINK_COUNT=$(nvidia-smi topo -m 2>/dev/null | grep -o 'NV[0-9][0-9]*' | wc -l) -if [ "$NVLINK_COUNT" -gt 0 ]; then - HAS_NVLINK=1 -else - HAS_NVLINK=0 -fi -echo "HAS_NVLINK: $HAS_NVLINK (detected $NVLINK_COUNT NVLink references)" - -SCRIPT_DIR="$(cd -- "$(dirname -- "${BASH_SOURCE[0]}")" &>/dev/null && pwd)" -source "${SCRIPT_DIR}/models/glm4.5-355B-A32B.sh" - -CKPT_ARGS=( - --hf-checkpoint $BASE_DIR/GLM-4.7-355B-A32B - --ref-load $BASE_DIR/GLM-4.7-355B-A32B_torch_dist/ -) - -ROLLOUT_ARGS=( - --prompt-data $BASE_DIR/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 128 - --n-samples-per-prompt 8 - --rollout-max-response-len 32768 - --rollout-temperature 1 - - --over-sampling-batch-size 256 - --dynamic-sampling-filter-path vime.rollout.filter_hub.dynamic_sampling_filters.check_reward_nonzero_std - - --num-steps-per-rollout 4 - --balance-data - --rollout-stop-token-ids 151329 151336 151338 -) - -EVAL_ARGS=( - --eval-interval 20 - --eval-prompt-data aime $BASE_DIR/rl_data/aime-2024.jsonl - --n-samples-per-eval-prompt 8 - --eval-max-response-len 32768 - --eval-top-p 1 -) - -PERF_ARGS=( - --tensor-model-parallel-size 8 - --sequence-parallel - --pipeline-model-parallel-size 4 - --context-parallel-size 2 - --expert-model-parallel-size 16 - --expert-tensor-parallel-size 1 - - --recompute-granularity full - --recompute-method uniform - --recompute-num-layers 1 - - --use-dynamic-batch-size - --max-tokens-per-gpu 16384 -) - -GRPO_ARGS=( - --advantage-estimator gspo - #--use-kl-loss - --kl-loss-coef 0.00 - --kl-loss-type low_var_kl - --kl-coef 0.00 - --entropy-coef 0.00 - --eps-clip 1e-4 - --eps-clip-high 2e-4 - - --use-tis -) - -OPTIMIZER_ARGS=( - --optimizer adam - --lr 1e-6 - --lr-decay-style constant - --weight-decay 0.1 - --adam-beta1 0.9 - --adam-beta2 0.98 - - --optimizer-cpu-offload - --overlap-cpu-optimizer-d2h-h2d - --use-precision-aware-optimizer -) - -WANDB_ARGS=( - # --use-wandb - # --wandb-project vime-dev - # --wandb-group glm4.7-355B -) - -VLLM_ARGS=( - --rollout-num-gpus-per-engine 32 - --vllm-gpu-memory-utilization 0.7 - --vllm-data-parallel-size 4 - - # mtp - --vllm-enable-expert-parallel - --vllm-speculative-config '{"method":"mtp","num_speculative_tokens":3}' -) - -MISC_ARGS=( - # default dropout in megatron is 0.1 - --attention-dropout 0.0 - --hidden-dropout 0.0 - # should be good for model performance - --accumulate-allreduce-grads-in-fp32 - --attention-softmax-in-fp32 - # need to comment this when using model with MLA - --attention-backend flash - - --moe-token-dispatcher-type flex - --moe-enable-deepep -) - -ACTOR_NUM_NODES=${ACTOR_NUM_NODES:-8} -ACTOR_NUM_GPUS_PER_NODE=${ACTOR_NUM_GPUS_PER_NODE:-8} -SOCKET_IFNAME=${SOCKET_IFNAME:-eth0} - -MASTER_ADDR=${MASTER_ADDR:-} -if [ -z "${MASTER_ADDR}" ]; then - echo "MASTER_ADDR is not set. Please set it to the master node address." - exit 1 -fi - -# launch the master node of ray -export no_proxy="127.0.0.1,${MASTER_ADDR}" -ray start --head --node-ip-address "${MASTER_ADDR}" --num-gpus "${ACTOR_NUM_GPUS_PER_NODE}" --disable-usage-stats --dashboard-host=0.0.0.0 --dashboard-port=8265 - -HOSTFILE=${HOSTFILE:-} -if [ -n "${HOSTFILE}" ]; then - for WORKER_IP in $(awk '{print $1}' "${HOSTFILE}"); do - if [[ "${WORKER_IP}" == "${MASTER_ADDR}" ]]; then - continue - fi - echo "Starting Ray worker on ${WORKER_IP}" - ssh root@"${WORKER_IP}" \ - "pkill -9 -f "vllm serve" ; ray stop --force ; pkill -9 python ; ray start --address=${MASTER_ADDR}:6379 --num-gpus ${ACTOR_NUM_GPUS_PER_NODE} --node-ip-address ${WORKER_IP} --disable-usage-stats" & - done - wait -fi - -RUNTIME_ENV_JSON=$(cat </dev/null | grep -o 'NV[0-9][0-9]*' | wc -l) -if [ "$NVLINK_COUNT" -gt 0 ]; then - HAS_NVLINK=1 -else - HAS_NVLINK=0 -fi -echo "HAS_NVLINK: $HAS_NVLINK (detected $NVLINK_COUNT NVLink references)" - -SCRIPT_DIR="$(cd -- "$(dirname -- "${BASH_SOURCE[0]}")" &>/dev/null && pwd)" -source "${SCRIPT_DIR}/models/glm5-744B-A40B.sh" - -CKPT_ARGS=( - --hf-checkpoint $BASE_DIR/GLM-5 - --ref-load $BASE_DIR/GLM-5_torch_dist/ - --load $BASE_DIR/GLM-5_slime/ - --save $BASE_DIR/GLM-5_slime/ - --save-interval 20 -) - -ROLLOUT_ARGS=( - --prompt-data $BASE_DIR/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 8 - --n-samples-per-prompt 8 - --rollout-max-response-len 32768 - --rollout-temperature 1 - - --global-batch-size 64 -) - -PERF_ARGS=( - --tensor-model-parallel-size 4 - --sequence-parallel - --pipeline-model-parallel-size 4 - --decoder-last-pipeline-num-layers 18 - --expert-model-parallel-size 32 - --expert-tensor-parallel-size 1 - --context-parallel-size 2 - - --recompute-granularity full - --recompute-method uniform - --recompute-num-layers 1 - - --use-dynamic-batch-size - --max-tokens-per-gpu 16384 - --data-pad-size-multiplier 4096 - --log-probs-chunk-size 1024 -) - -GRPO_ARGS=( - --advantage-estimator grpo - #--use-kl-loss - --kl-loss-coef 0.00 - --kl-loss-type low_var_kl - --kl-coef 0.00 - --entropy-coef 0.00 - --eps-clip 0.2 - --eps-clip-high 0.28 -) - -OPTIMIZER_ARGS=( - --optimizer adam - --lr 1e-6 - - --lr-decay-style constant - --weight-decay 0.1 - --adam-beta1 0.9 - --adam-beta2 0.98 - - --optimizer-cpu-offload - --overlap-cpu-optimizer-d2h-h2d - --use-precision-aware-optimizer -) - -WANDB_ARGS=( - # --use-wandb - # --wandb-project vime-dev - # --wandb-group glm5-test - # --wandb-key ${WANDB_KEY} -) - -VLLM_ARGS=( - --rollout-num-gpus-per-engine 64 - --vllm-gpu-memory-utilization 0.70 - --vllm-data-parallel-size 64 - - - --prefill-num-servers 1 - - # mtp - - # dsa - --vllm-attention-backend FLASHMLA_SPARSE - - --vllm-max-num-seqs 512 - --vllm-enable-expert-parallel - --vllm-all2all-backend deepep_high_throughput - --vllm-speculative-config '{"method":"mtp","num_speculative_tokens":3}' - --vllm-block-size 64 - --vllm-max-cudagraph-capture-size 8 - --vllm-max-num-batched-tokens 131072 -) - -MISC_ARGS=( - # default dropout in megatron is 0.1 - --attention-dropout 0.0 - --hidden-dropout 0.0 - # should be good for model performance - --accumulate-allreduce-grads-in-fp32 - --attention-softmax-in-fp32 - # need to comment this when using model with MLA - --attention-backend flash - - # use deepep for megatron - --moe-enable-deepep - --moe-token-dispatcher-type flex -) - -# Build the runtime environment JSON with proper variable substitution -RUNTIME_ENV_JSON="{ - \"env_vars\": { - \"PYTHONPATH\": \"/root/Megatron-LM/\", - \"CUDA_DEVICE_MAX_CONNECTIONS\": \"1\", - \"NCCL_NVLS_ENABLE\": \"${HAS_NVLINK}\", - \"no_proxy\": \"${no_proxy}\", - \"MASTER_ADDR\": \"${MASTER_ADDR}\", - \"INDEXER_ROPE_NEOX_STYLE\": \"0\", - \"NVSHMEM_DISABLE_NCCL\": \"1\" - } -}" - -ray job submit --address="http://127.0.0.1:8265" \ - --runtime-env-json="${RUNTIME_ENV_JSON}" \ --- python3 train.py \ - --actor-num-nodes 32 \ - --actor-num-gpus-per-node 8 \ - --colocate \ - --update-weight-buffer-size $(( 1024 * 1024 * 1024 * 2 )) \ - ${MODEL_ARGS[@]} \ - ${CKPT_ARGS[@]} \ - ${ROLLOUT_ARGS[@]} \ - ${OPTIMIZER_ARGS[@]} \ - ${GRPO_ARGS[@]} \ - ${WANDB_ARGS[@]} \ - ${PERF_ARGS[@]} \ - ${EVAL_ARGS[@]} \ - ${VLLM_ARGS[@]} \ - ${MISC_ARGS[@]} diff --git a/scripts/run-gpt-oss-20B.sh b/scripts/run-gpt-oss-20B.sh deleted file mode 100755 index 8668bc619..000000000 --- a/scripts/run-gpt-oss-20B.sh +++ /dev/null @@ -1,160 +0,0 @@ -#!/bin/bash - -# GPT-OSS 20B training script — single-node 8×H100 -# Model: openai/gpt-oss-20b (20B MoE, 32 experts top-4) -# -# Prerequisites: -# 1. Preprocess: python tools/preprocess_gpt_oss.py --input /path/to/gpt-oss-20b --output /path/to/gpt-oss-20b-bf16 -# 2. Convert: torchrun --nproc_per_node 8 tools/convert_hf_to_torch_dist.py \ -# --hf-checkpoint /path/to/gpt-oss-20b-bf16 \ -# --save /path/to/gpt-oss-20b_torch_dist \ -# --megatron-to-hf-mode bridge \ -# $(cat scripts/models/gpt-oss-20B.sh | grep -oP "'[^']*'|--[^ ]+( [^ -][^ ]*)?") - -# for rerun the task -pkill -9 -f "vllm serve" -sleep 3 -ray stop --force -pkill -9 ray -pkill -9 python -sleep 3 -pkill -9 ray -pkill -9 python - -set -ex - -export PYTHONBUFFERED=16 - -NVLINK_COUNT=$(nvidia-smi topo -m 2>/dev/null | grep -o 'NV[0-9][0-9]*' | wc -l) -if [ "$NVLINK_COUNT" -gt 0 ]; then - HAS_NVLINK=1 -else - HAS_NVLINK=0 -fi -echo "HAS_NVLINK: $HAS_NVLINK (detected $NVLINK_COUNT NVLink references)" - -SCRIPT_DIR="$(cd -- "$(dirname -- "${BASH_SOURCE[0]}")" &>/dev/null && pwd)" -source "${SCRIPT_DIR}/models/gpt-oss-20B.sh" - -CKPT_ARGS=( - --hf-checkpoint /root/gpt-oss-20b-bf16 - --ref-load /root/gpt-oss-20b_torch_dist - --load /root/gpt-oss-20b_slime/ - --save /root/gpt-oss-20b_slime/ - --save-interval 20 -) - -ROLLOUT_ARGS=( - --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_ARGS=( - --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 -) - -PERF_ARGS=( - --tensor-model-parallel-size 1 - --pipeline-model-parallel-size 2 - --context-parallel-size 1 - --expert-model-parallel-size 4 - --expert-tensor-parallel-size 1 - - --recompute-granularity full - --recompute-method uniform - --recompute-num-layers 1 - - --use-dynamic-batch-size - --max-tokens-per-gpu 1536 - # PP>1 requires seq-length override (PP recv buffer is sized by seq-length). - # Must be >= max-tokens-per-gpu. Smaller = less memory for FP32 logits. - --seq-length 1536 - # Reduce memory margin to avoid excessive CPU↔GPU swapping in colocate mode. - --train-memory-margin-bytes 268435456 -) - -GRPO_ARGS=( - --advantage-estimator grpo - --use-kl-loss - --kl-loss-coef 0.001 - --num-steps-per-rollout 1 -) - -OPTIMIZER_ARGS=( - --lr 5e-7 - --lr-decay-style cosine - --min-lr 0 - --lr-warmup-fraction 0.01 - --weight-decay 0.01 - --adam-beta1 0.9 - --adam-beta2 0.99 - --clip-grad 1.0 - --micro-batch-size 1 -) - -WANDB_ARGS=( - # --wandb-project gpt-oss-20b - # --wandb-exp-name gpt-oss-20b-grpo -) - -VLLM_ARGS=( - --rollout-num-gpus 8 - --vllm-tensor-parallel-size 1 - --vllm-gpu-memory-utilization 0.55 - --vllm-max-num-seqs 64 - --vllm-max-cudagraph-capture-size 16 -) - -MISC_ARGS=( - --attention-dropout 0.0 - --hidden-dropout 0.0 - --accumulate-allreduce-grads-in-fp32 - --attention-softmax-in-fp32 - --moe-token-dispatcher-type alltoall - --megatron-to-hf-mode bridge -) - -# launch the master node of ray in container -export MASTER_ADDR=${MASTER_ADDR:-"127.0.0.1"} -ray start --head --node-ip-address ${MASTER_ADDR} --num-gpus 8 --disable-usage-stats --dashboard-host=0.0.0.0 --dashboard-port=8265 - -RUNTIME_ENV_JSON="{ - \"env_vars\": { - \"PYTHONPATH\": \"/root/Megatron-LM/\", - \"CUDA_DEVICE_MAX_CONNECTIONS\": \"1\", - \"NCCL_NVLS_ENABLE\": \"${HAS_NVLINK}\" - } -}" - -ray job submit --address="http://127.0.0.1:8265" \ - --runtime-env-json="${RUNTIME_ENV_JSON}" \ - -- python3 train.py \ - --actor-num-nodes 1 \ - --actor-num-gpus-per-node 8 \ - --colocate \ - ${MODEL_ARGS[@]} \ - ${CKPT_ARGS[@]} \ - ${ROLLOUT_ARGS[@]} \ - ${OPTIMIZER_ARGS[@]} \ - ${GRPO_ARGS[@]} \ - ${WANDB_ARGS[@]} \ - ${PERF_ARGS[@]} \ - ${EVAL_ARGS[@]} \ - ${VLLM_ARGS[@]} \ - ${MISC_ARGS[@]} diff --git a/scripts/run-kimi-k2-Instruct.sh b/scripts/run-kimi-k2-Instruct.sh deleted file mode 100644 index bfb463dd0..000000000 --- a/scripts/run-kimi-k2-Instruct.sh +++ /dev/null @@ -1,176 +0,0 @@ -#!/bin/bash - -# for rerun the task -pkill -9 -f "vllm serve" -sleep 3 -ray stop --force -pkill -9 ray -pkill -9 python -sleep 3 -pkill -9 ray -pkill -9 python - -set -ex - -# will prevent ray from buffering stdout/stderr -export PYTHONBUFFERED=16 - -NVLINK_COUNT=$(nvidia-smi topo -m 2>/dev/null | grep -o 'NV[0-9][0-9]*' | wc -l) -if [ "$NVLINK_COUNT" -gt 0 ]; then - HAS_NVLINK=1 -else - HAS_NVLINK=0 -fi -echo "HAS_NVLINK: $HAS_NVLINK (detected $NVLINK_COUNT NVLink references)" - -SCRIPT_DIR="$(cd -- "$(dirname -- "${BASH_SOURCE[0]}")" &>/dev/null && pwd)" -source "${SCRIPT_DIR}/models/kimi-k2.sh" - -CKPT_ARGS=( - --hf-checkpoint $BASE_DIR/Kimi-K2-Instruct/ - # --hf-checkpoint $BASE_DIR/Kimi-K2-bf16/ - --ref-load $BASE_DIR/Kimi-K2_torch_dist/ - --load $BASE_DIR/Kimi-K2_slime/ - --save $BASE_DIR/Kimi-K2_slime/ - --save-interval 20 -) - -ROLLOUT_ARGS=( - --prompt-data $BASE_DIR/dapo-math-17k/dapo-math-17k.jsonl - --input-key prompt - --label-key label - --apply-chat-template - --rollout-shuffle - - --rm-type math - - --num-rollout 100 - --rollout-batch-size 128 - --n-samples-per-prompt 8 - --rollout-max-response-len 32768 - --rollout-temperature 1 - - # --global-batch-size 1024 - - --over-sampling-batch-size 256 - --dynamic-sampling-filter-path vime.rollout.filter_hub.dynamic_sampling_filters.check_reward_nonzero_std - - --num-steps-per-rollout 4 - --balance-data -) - -EVAL_ARGS=( - --eval-interval 20 - --eval-prompt-data aime $BASE_DIR/rl_data/aime-2024.jsonl - --n-samples-per-eval-prompt 8 - --eval-max-response-len 32768 - --eval-top-p 1 -) - -PERF_ARGS=( - --tensor-model-parallel-size 8 - --sequence-parallel - --pipeline-model-parallel-size 8 - --context-parallel-size 4 - --expert-model-parallel-size 32 - --expert-tensor-parallel-size 1 - --decoder-last-pipeline-num-layers 5 - - --recompute-granularity full - --recompute-method uniform - --recompute-num-layers 1 - - --use-dynamic-batch-size - --max-tokens-per-gpu 16384 -) - -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 -) - -OPTIMIZER_ARGS=( - --optimizer adam - --lr 1e-6 - - --lr-decay-style constant - --weight-decay 0.1 - --adam-beta1 0.9 - --adam-beta2 0.98 - - --optimizer-cpu-offload - --overlap-cpu-optimizer-d2h-h2d - --use-precision-aware-optimizer -) - -WANDB_ARGS=( - # --use-wandb - # --wandb-project vime-dev - # --wandb-group kimi-k2-test - # --wandb-key ${WANDB_KEY} -) - -VLLM_ARGS=( - --rollout-num-gpus-per-engine 16 - --vllm-gpu-memory-utilization 0.7 - - # dp attention - --vllm-data-parallel-size 8 - - - # deepep MoE A2A — add --vllm-all2all-backend deepep_high_throughput to enable - - # make every dp rank has 128 concurrency - --vllm-server-concurrency 1024 - --vllm-enable-expert-parallel -) - - -MISC_ARGS=( - # default dropout in megatron is 0.1 - --attention-dropout 0.0 - --hidden-dropout 0.0 - # should be good for model performance - --accumulate-allreduce-grads-in-fp32 - --attention-softmax-in-fp32 - # need to comment this when using model with MLA - --attention-backend flash - - # use deepep for megatron - --moe-enable-deepep - --moe-token-dispatcher-type flex -) - -# Build the runtime environment JSON with proper variable substitution -RUNTIME_ENV_JSON="{ - \"env_vars\": { - \"PYTHONPATH\": \"/root/Megatron-LM/\", - \"CUDA_DEVICE_MAX_CONNECTIONS\": \"1\", - \"NCCL_NVLS_ENABLE\": \"${HAS_NVLINK}\", - \"no_proxy\": \"${no_proxy}\", - \"MASTER_ADDR\": \"${MASTER_ADDR}\" - } -}" - -ray job submit --address="http://127.0.0.1:8265" \ - --runtime-env-json="${RUNTIME_ENV_JSON}" \ - -- python3 train.py \ - --actor-num-nodes 32 \ - --actor-num-gpus-per-node 8 \ - --colocate \ - --update-weight-buffer-size $(( 4 * 512 * 1024 * 1024)) - ${MODEL_ARGS[@]} \ - ${CKPT_ARGS[@]} \ - ${ROLLOUT_ARGS[@]} \ - ${OPTIMIZER_ARGS[@]} \ - ${GRPO_ARGS[@]} \ - ${WANDB_ARGS[@]} \ - ${PERF_ARGS[@]} \ - ${EVAL_ARGS[@]} \ - ${VLLM_ARGS[@]} \ - ${MISC_ARGS[@]} diff --git a/scripts/run-kimi-k2-Thinking.sh b/scripts/run-kimi-k2-Thinking.sh deleted file mode 100644 index 32474350a..000000000 --- a/scripts/run-kimi-k2-Thinking.sh +++ /dev/null @@ -1,178 +0,0 @@ -#!/bin/bash - -# for rerun the task -pkill -9 -f "vllm serve" -sleep 3 -ray stop --force -pkill -9 ray -pkill -9 python -sleep 3 -pkill -9 ray -pkill -9 python - -set -ex - -# will prevent ray from buffering stdout/stderr -export PYTHONBUFFERED=16 - -NVLINK_COUNT=$(nvidia-smi topo -m 2>/dev/null | grep -o 'NV[0-9][0-9]*' | wc -l) -if [ "$NVLINK_COUNT" -gt 0 ]; then - HAS_NVLINK=1 -else - HAS_NVLINK=0 -fi -echo "HAS_NVLINK: $HAS_NVLINK (detected $NVLINK_COUNT NVLink references)" - -SCRIPT_DIR="$(cd -- "$(dirname -- "${BASH_SOURCE[0]}")" &>/dev/null && pwd)" -source "${SCRIPT_DIR}/models/kimi-k2-thinking.sh" - -CKPT_ARGS=( - # --hf-checkpoint $BASE_DIR/Kimi-K2-Thinking-bf16/ - --hf-checkpoint $BASE_DIR/Kimi-K2-Thinking-fp8/ - --ref-load $BASE_DIR/Kimi-K2-Thinking_torch_dist/ - --load $BASE_DIR/Kimi-K2-Thinking_slime/ - --save $BASE_DIR/Kimi-K2-Thinking_slime/ - --save-interval 20 -) - -ROLLOUT_ARGS=( - --prompt-data $BASE_DIR/dapo-math-17k/dapo-math-17k.jsonl - --input-key prompt - --label-key label - --apply-chat-template - --rollout-shuffle - - --rm-type math - - --num-rollout 100 - --rollout-batch-size 128 - --n-samples-per-prompt 8 - --rollout-max-response-len 16384 - --rollout-temperature 1 - - # --global-batch-size 1024 - - --over-sampling-batch-size 256 - --dynamic-sampling-filter-path vime.rollout.filter_hub.dynamic_sampling_filters.check_reward_nonzero_std - - --num-steps-per-rollout 4 - --balance-data -) - -EVAL_ARGS=( - --eval-interval 20 - --eval-prompt-data aime $BASE_DIR/aime-2024.jsonl - --n-samples-per-eval-prompt 16 - --eval-max-response-len 16384 - --eval-top-p 1 -) - -PERF_ARGS=( - --tensor-model-parallel-size 8 - --sequence-parallel - --pipeline-model-parallel-size 8 - --context-parallel-size 4 - --expert-model-parallel-size 32 - --expert-tensor-parallel-size 1 - --decoder-last-pipeline-num-layers 5 - - --recompute-granularity full - --recompute-method uniform - --recompute-num-layers 1 - - --use-dynamic-batch-size - --max-tokens-per-gpu 16384 -) - -GRPO_ARGS=( - --advantage-estimator grpo - --use-kl-loss - --kl-loss-coef 0.00 - --kl-loss-type low_var_kl - # --kl-coef 0.00 - --entropy-coef 0.00 - --eps-clip 0.2 - --eps-clip-high 0.28 - --use-tis -) - -OPTIMIZER_ARGS=( - --optimizer adam - --lr 1e-6 - - --lr-decay-style constant - --weight-decay 0.1 - --adam-beta1 0.9 - --adam-beta2 0.98 - - --optimizer-cpu-offload - --overlap-cpu-optimizer-d2h-h2d - --use-precision-aware-optimizer -) - -WANDB_ARGS=( - # --use-wandb - # --wandb-project vime-dev - # --wandb-group kimi-k2-test - # --wandb-key ${WANDB_KEY} -) - -VLLM_ARGS=( - --rollout-num-gpus-per-engine 16 - --vllm-gpu-memory-utilization 0.7 - - # dp attention - --vllm-data-parallel-size 8 - - - # deepep MoE A2A — add --vllm-all2all-backend deepep_high_throughput to enable - - # make every dp rank has 128 concurrency - --vllm-server-concurrency 1024 - --vllm-enable-expert-parallel -) - - -MISC_ARGS=( - # default dropout in megatron is 0.1 - --attention-dropout 0.0 - --hidden-dropout 0.0 - # should be good for model performance - --accumulate-allreduce-grads-in-fp32 - --attention-softmax-in-fp32 - # need to comment this when using model with MLA - --attention-backend flash - - # use deepep for megatron - # --moe-enable-deepep - # --moe-token-dispatcher-type flex -) - -# Build the runtime environment JSON with proper variable substitution -RUNTIME_ENV_JSON="{ - \"env_vars\": { - \"PYTHONPATH\": \"/root/Megatron-LM/\", - \"CUDA_DEVICE_MAX_CONNECTIONS\": \"1\", - \"NCCL_NVLS_ENABLE\": \"${HAS_NVLINK}\", - \"no_proxy\": \"${no_proxy}\", - \"MASTER_ADDR\": \"${MASTER_ADDR}\" - } -}" - -ray job submit --address="http://127.0.0.1:8265" \ - --runtime-env-json="${RUNTIME_ENV_JSON}" \ - -- python3 train.py \ - --actor-num-nodes 32 \ - --actor-num-gpus-per-node 8 \ - --colocate \ - --update-weight-buffer-size $(( 4 * 512 * 1024 * 1024)) \ - ${MODEL_ARGS[@]} \ - ${CKPT_ARGS[@]} \ - ${ROLLOUT_ARGS[@]} \ - ${OPTIMIZER_ARGS[@]} \ - ${GRPO_ARGS[@]} \ - ${WANDB_ARGS[@]} \ - ${PERF_ARGS[@]} \ - ${EVAL_ARGS[@]} \ - ${VLLM_ARGS[@]} \ - ${MISC_ARGS[@]} diff --git a/scripts/run-mimo-7B-rl-eagle.sh b/scripts/run-mimo-7B-rl-eagle.sh deleted file mode 100644 index 4bbee9b38..000000000 --- a/scripts/run-mimo-7B-rl-eagle.sh +++ /dev/null @@ -1,163 +0,0 @@ - -#!/bin/bash - -# for rerun the task -pkill -9 -f "vllm serve" -sleep 3 -ray stop --force -pkill -9 ray -pkill -9 python -sleep 3 -pkill -9 ray -pkill -9 python - -set -ex - -# will prevent ray from buffering stdout/stderr -export PYTHONBUFFERED=16 - -NVLINK_COUNT=$(nvidia-smi topo -m 2>/dev/null | grep -o 'NV[0-9][0-9]*' | wc -l) -if [ "$NVLINK_COUNT" -gt 0 ]; then - HAS_NVLINK=1 -else - HAS_NVLINK=0 -fi -echo "HAS_NVLINK: $HAS_NVLINK (detected $NVLINK_COUNT NVLink references)" - -SCRIPT_DIR="$(cd -- "$(dirname -- "${BASH_SOURCE[0]}")" &>/dev/null && pwd)" -source "${SCRIPT_DIR}/models/mimo-7B-rl.sh" - - CKPT_ARGS=( - --hf-checkpoint /root/MiMo-7B-RL - #--hf-checkpoint /root/Qwen3-4B-FP8 - --ref-load /root/MiMo-7B-RL_torch_dist - --load /root/MiMo-7B-RL-mtp_slime/ - --save /root/MiMo-7B-RL-mtp_slime/ - --save-interval 2000 -) - -ROLLOUT_ARGS=( - --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_ARGS=( - --eval-interval 20 - --eval-prompt-data aime /root/aime-2024/aime-2024.jsonl - --n-samples-per-eval-prompt 1 - --eval-max-response-len 8192 - --eval-top-p 1 -) - -PERF_ARGS=( - --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 - - # --micro-batch-size 1 - --use-dynamic-batch-size - --max-tokens-per-gpu 9216 -) - -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 -) - -OPTIMIZER_ARGS=( - --optimizer adam - --lr 1e-6 - --lr-decay-style constant - --weight-decay 0.1 - --adam-beta1 0.9 - --adam-beta2 0.98 -) - -WANDB_ARGS=( - # --use-wandb - # --wandb-project vime-dev - # --wandb-group mimo-7B-rl-test - # --wandb-key ${WANDB_API_KEY} -) - -VLLM_ARGS=( - --rollout-num-gpus-per-engine 1 - --vllm-gpu-memory-utilization 0.7 - - # for speculative decoding - - # sometimes flashinfer has IMA bugs. Use fa3 as instead - --vllm-attention-backend FLASH_ATTN - --vllm-speculative-config '{"method":"mtp","num_speculative_tokens":3}' -) - -MISC_ARGS=( - # default dropout in megatron is 0.1 - --attention-dropout 0.0 - --hidden-dropout 0.0 - # should be good for model performance - --accumulate-allreduce-grads-in-fp32 - --attention-softmax-in-fp32 - # need to comment this when using model with MLA - --attention-backend flash -) - -SPEC_ARGS=( - --enable-mtp-training - --mtp-loss-scaling-factor 0.2 -) - -# launch the master node of ray in container -export MASTER_ADDR=${MASTER_ADDR:-"127.0.0.1"} -ray start --head --node-ip-address ${MASTER_ADDR} --num-gpus 8 --disable-usage-stats - -# Build the runtime environment JSON with proper variable substitution -RUNTIME_ENV_JSON="{ - \"env_vars\": { - \"PYTHONPATH\": \"/root/Megatron-LM/\", - \"CUDA_DEVICE_MAX_CONNECTIONS\": \"1\", - \"NCCL_NVLS_ENABLE\": \"${HAS_NVLINK}\" - } -}" - -ray job submit --address="http://127.0.0.1:8265" \ - --runtime-env-json="${RUNTIME_ENV_JSON}" \ - -- python3 train.py \ - --actor-num-nodes 1 \ - --actor-num-gpus-per-node 8 \ - --colocate \ - ${MODEL_ARGS[@]} \ - ${CKPT_ARGS[@]} \ - ${ROLLOUT_ARGS[@]} \ - ${OPTIMIZER_ARGS[@]} \ - ${GRPO_ARGS[@]} \ - ${WANDB_ARGS[@]} \ - ${PERF_ARGS[@]} \ - ${EVAL_ARGS[@]} \ - ${VLLM_ARGS[@]} \ - ${MISC_ARGS[@]} \ - ${SPEC_ARGS[@]} diff --git a/scripts/run-moonlight-16B-A3B.sh b/scripts/run-moonlight-16B-A3B.sh deleted file mode 100644 index 02ee97916..000000000 --- a/scripts/run-moonlight-16B-A3B.sh +++ /dev/null @@ -1,163 +0,0 @@ -#!/bin/bash - -# for rerun the task -pkill -9 -f "vllm serve" -sleep 3 -ray stop --force -pkill -9 ray -pkill -9 python -sleep 3 -pkill -9 ray -pkill -9 python -pkill -9 redis - -set -ex - -# will prevent ray from buffering stdout/stderr -export PYTHONBUFFERED=16 - -NVLINK_COUNT=$(nvidia-smi topo -m 2>/dev/null | grep -o 'NV[0-9][0-9]*' | wc -l) -if [ "$NVLINK_COUNT" -gt 0 ]; then - HAS_NVLINK=1 -else - HAS_NVLINK=0 -fi -echo "HAS_NVLINK: $HAS_NVLINK (detected $NVLINK_COUNT NVLink references)" - -SCRIPT_DIR="$(cd -- "$(dirname -- "${BASH_SOURCE[0]}")" &>/dev/null && pwd)" -source "${SCRIPT_DIR}/models/moonlight.sh" - -CKPT_ARGS=( - --hf-checkpoint /root/Moonlight-16B-A3B - --ref-load /root/Moonlight-16B-A3B_torch_dist - --load /root/Moonlight-16B-A3B_slime/ - --save /root/Moonlight-16B-A3B_slime/ - --save-interval 20 -) - -ROLLOUT_ARGS=( - --prompt-data /root/dapo-math-17k/dapo-math-17k.jsonl - --input-key prompt - --label-key label - --apply-chat-template - --rollout-shuffle - --rm-type math - --num-rollout 3000 - --rollout-batch-size 128 - --n-samples-per-prompt 8 - --rollout-max-response-len 4096 - --rollout-temperature 1 - - --over-sampling-batch-size 256 - --dynamic-sampling-filter-path vime.rollout.filter_hub.dynamic_sampling_filters.check_reward_nonzero_std - - --num-steps-per-rollout 4 - # --global-batch-size 256 - --balance-data -) - -EVAL_ARGS=( - --eval-interval 20 - --eval-prompt-data aime /root/aime-2024/aime-2024.jsonl - --n-samples-per-eval-prompt 8 - --eval-max-response-len 4096 - --eval-top-p 1 -) - -PERF_ARGS=( - --tensor-model-parallel-size 4 - --sequence-parallel - --pipeline-model-parallel-size 1 - --context-parallel-size 1 - --expert-model-parallel-size 8 - --expert-tensor-parallel-size 1 - - --recompute-granularity full - --recompute-method uniform - --recompute-num-layers 1 - - # --micro-batch-size 1 - --use-dynamic-batch-size - --max-tokens-per-gpu 8192 -) - -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 -) - -OPTIMIZER_ARGS=( - --optimizer adam - --lr 1e-6 - --lr-decay-style constant - --weight-decay 0.1 - --adam-beta1 0.9 - --adam-beta2 0.98 - - --optimizer-cpu-offload - --overlap-cpu-optimizer-d2h-h2d - --use-precision-aware-optimizer -) - -WANDB_ARGS=( - # --use-wandb - # --wandb-project vime-dev - # --wandb-group moomlight-16B-A3B-test - # --wandb-key ${WANDB_KEY} -) - -VLLM_ARGS=( - --rollout-num-gpus-per-engine 8 - --vllm-gpu-memory-utilization 0.7 - --vllm-cudagraph-capture-sizes 1 2 4 8 $(seq 16 8 256) -) - -MISC_ARGS=( - # default dropout in megatron is 0.1 - --attention-dropout 0.0 - --hidden-dropout 0.0 - # should be good for model performance - --accumulate-allreduce-grads-in-fp32 - --attention-softmax-in-fp32 - # need to comment this when using model with MLA - # --attention-backend flash - - # use deepep for megatron - --moe-enable-deepep - --moe-token-dispatcher-type flex -) - -# launch the master node of ray in container -export MASTER_ADDR=${MASTER_ADDR:-"127.0.0.1"} -ray start --head --node-ip-address ${MASTER_ADDR} --num-gpus 8 --disable-usage-stats --dashboard-host=0.0.0.0 --dashboard-port=8265 - -# Build the runtime environment JSON with proper variable substitution -RUNTIME_ENV_JSON="{ - \"env_vars\": { - \"PYTHONPATH\": \"/root/Megatron-LM/\", - \"CUDA_DEVICE_MAX_CONNECTIONS\": \"1\", - \"NCCL_NVLS_ENABLE\": \"${HAS_NVLINK}\" - } -}" - -ray job submit --address="http://127.0.0.1:8265" \ - --runtime-env-json="${RUNTIME_ENV_JSON}" \ - -- python3 train.py \ - --actor-num-nodes 1 \ - --actor-num-gpus-per-node 8 \ - --colocate \ - ${MODEL_ARGS[@]} \ - ${CKPT_ARGS[@]} \ - ${ROLLOUT_ARGS[@]} \ - ${OPTIMIZER_ARGS[@]} \ - ${GRPO_ARGS[@]} \ - ${WANDB_ARGS[@]} \ - ${PERF_ARGS[@]} \ - ${EVAL_ARGS[@]} \ - ${VLLM_ARGS[@]} \ - ${MISC_ARGS[@]} diff --git a/scripts/run-qwen2.5-0.5B-reproducibility.sh b/scripts/run-qwen2.5-0.5B-reproducibility.sh deleted file mode 100644 index e8d48e1af..000000000 --- a/scripts/run-qwen2.5-0.5B-reproducibility.sh +++ /dev/null @@ -1,138 +0,0 @@ -#!/bin/bash - -# for rerun the task -pkill -9 -f "vllm serve" -sleep 3 -ray stop --force -pkill -9 ray -pkill -9 python -sleep 3 -pkill -9 ray -pkill -9 python - -set -ex - -# will prevent ray from buffering stdout/stderr -export PYTHONBUFFERED=16 - -SCRIPT_DIR="$(cd -- "$(dirname -- "${BASH_SOURCE[0]}")" &>/dev/null && pwd)" -source "${SCRIPT_DIR}/scripts/models/qwen2.5-0.5B.sh" - -CKPT_ARGS=( - --hf-checkpoint /root/Qwen2.5-0.5B-Instruct/ - --ref-load /root/Qwen2.5-0.5B-Instruct_torch_dist/ -) - -ROLLOUT_ARGS=( - --prompt-data /root/gsm8k/train.parquet - --input-key messages - --label-key label - --apply-chat-template - --rollout-shuffle - --rm-type math - --num-rollout 100 - --rollout-batch-size 32 - --n-samples-per-prompt 8 - --rollout-max-response-len 1024 - --rollout-temperature 1 - - --global-batch-size 256 -) - -EVAL_ARGS=( - --eval-interval 20 - --eval-prompt-data gsm8k /root/gsm8k/test.parquet - --n-samples-per-eval-prompt 1 - --eval-max-response-len 1024 - --eval-top-k 1 -) - -PERF_ARGS=( - --tensor-model-parallel-size 1 - --sequence-parallel - --pipeline-model-parallel-size 1 - --context-parallel-size 1 - --expert-model-parallel-size 1 - --expert-tensor-parallel-size 1 - - --use-dynamic-batch-size - --max-tokens-per-gpu 9216 -) - -GRPO_ARGS=( - --advantage-estimator grpo - --use-kl-loss - --kl-loss-coef 0.00 - --kl-loss-type low_var_kl - --kl-coef 0.00 - --entropy-coef 0.00 - --eps-clip 0.2 - --eps-clip-high 0.28 -) - -OPTIMIZER_ARGS=( - --optimizer adam - --lr 1e-6 - --lr-decay-style constant - --weight-decay 0.1 - --adam-beta1 0.9 - --adam-beta2 0.98 -) - -WANDB_ARGS=( - --use-wandb - --wandb-host https://wandb.ai/ - --wandb-team glm-zero - --wandb-project vime-dev - --wandb-group qwen2.5-0.5B-gsm8k-deterministic -) - -VLLM_ARGS=( - --rollout-num-gpus-per-engine 1 - --vllm-gpu-memory-utilization 0.7 - - --vllm-attention-backend FLASHINFER - - --deterministic-mode -) - -MISC_ARGS=( - # default dropout in megatron is 0.1 - --attention-dropout 0.0 - --hidden-dropout 0.0 - # should be good for model performance - --accumulate-allreduce-grads-in-fp32 - --attention-softmax-in-fp32 - # need to comment this when using model with MLA - --attention-backend flash -) - -# launch the master node of ray in container -ray start --head --node-ip-address 127.0.0.1 --num-gpus 8 --disable-usage-stats - -ray job submit --address="http://127.0.0.1:8265" \ - --runtime-env-json='{ - "env_vars": { - "PYTHONPATH": "/root/Megatron-LM", - "CUDA_DEVICE_MAX_CONNECTIONS": "1", - "NCCL_ALGO": "Ring", - "NVTE_ALLOW_NONDETERMINISTIC_ALGO": "0", - "CUBLAS_WORKSPACE_CONFIG": ":4096:8", - "VLLM_BATCH_INVARIANT": "1" - } - }' \ - -- python3 train.py \ - --actor-num-nodes 1 \ - --actor-num-gpus-per-node 8 \ - --colocate \ - --calculate-per-token-loss \ - ${MODEL_ARGS[@]} \ - ${CKPT_ARGS[@]} \ - ${ROLLOUT_ARGS[@]} \ - ${OPTIMIZER_ARGS[@]} \ - ${GRPO_ARGS[@]} \ - ${WANDB_ARGS[@]} \ - ${PERF_ARGS[@]} \ - ${EVAL_ARGS[@]} \ - ${VLLM_ARGS[@]} \ - ${MISC_ARGS[@]} diff --git a/scripts/run-qwen3-235B-A22B-sft.sh b/scripts/run-qwen3-235B-A22B-sft.sh deleted file mode 100644 index 6b75af8fc..000000000 --- a/scripts/run-qwen3-235B-A22B-sft.sh +++ /dev/null @@ -1,151 +0,0 @@ -#!/bin/bash - -# for rerun the task -pkill -9 -f "vllm serve" -sleep 3 -ray stop --force -pkill -9 ray -pkill -9 python -sleep 3 -pkill -9 ray -pkill -9 python - -set -ex - -# if base folder not set raise error -if [ -z "${BASE_FOLDER}" ]; then - echo "BASE_FOLDER is not set. Please set it to the base directory of your checkpoints." - exit 1 -fi - -if [ -z "${MASTER_ADDR}" ]; then - echo "MASTER_ADDR is not set. Please set it to the master node address." - exit 1 -fi - -# will prevent ray from buffering stdout/stderr -export PYTHONBUFFERED=16 - -NVLINK_COUNT=$(nvidia-smi topo -m 2>/dev/null | grep -o 'NV[0-9][0-9]*' | wc -l) -if [ "$NVLINK_COUNT" -gt 0 ]; then - HAS_NVLINK=1 -else - HAS_NVLINK=0 -fi -echo "HAS_NVLINK: $HAS_NVLINK (detected $NVLINK_COUNT NVLink references)" - -SCRIPT_DIR="$(cd -- "$(dirname -- "${BASH_SOURCE[0]}")" &>/dev/null && pwd)" -source "${SCRIPT_DIR}/models/qwen3-235B-A22B.sh" - -CKPT_ARGS=( - --hf-checkpoint ${BASE_FOLDER}/Qwen3-235B-A22B - --ref-load ${BASE_FOLDER}/Qwen3-235B-A22B_torch_dist - --load ${BASE_FOLDER}/Qwen3-235B-A22B_slime/ - --save ${BASE_FOLDER}/Qwen3-235B-A22B_slime/ - --save-interval 1000 -) - -SFT_ARGS=( - --rollout-function-path vime.rollout.sft_rollout.generate_rollout - --prompt-data ${BASE_FOLDER}/openhermes2_5.parquet - --input-key messages - # --apply-chat-template - --rollout-shuffle - --num-epoch 3 - --rollout-batch-size 128 - --global-batch-size 128 - - --loss-type sft_loss - --calculate-per-token-loss - --disable-compute-advantages-and-returns - --debug-train-only -) - -PERF_ARGS=( - --tensor-model-parallel-size 4 - --sequence-parallel - --pipeline-model-parallel-size 1 - --context-parallel-size 1 - --expert-model-parallel-size 32 - --expert-tensor-parallel-size 1 - - --recompute-granularity full - --recompute-method uniform - --recompute-num-layers 1 - - # --micro-batch-size 1 - --use-dynamic-batch-size - --max-tokens-per-gpu 9216 -) - -OPTIMIZER_ARGS=( - --optimizer adam - --lr 1e-5 - --lr-decay-style cosine - --min-lr 1e-6 - --lr-warmup-fraction 0.1 - --weight-decay 0.1 - --adam-beta1 0.9 - --adam-beta2 0.98 - - --optimizer-cpu-offload - --overlap-cpu-optimizer-d2h-h2d - --use-precision-aware-optimizer -) - -WANDB_ARGS=( - # --use-wandb - # --wandb-project vime-dev - # --wandb-group qwen3-235B-sft -) - -MISC_ARGS=( - # default dropout in megatron is 0.1 - --attention-dropout 0.0 - --hidden-dropout 0.0 - # should be good for model performance - --accumulate-allreduce-grads-in-fp32 - --attention-softmax-in-fp32 - # need to comment this when using model with MLA - --attention-backend flash -) - -# launch the master node of ray in container -export no_proxy="127.0.0.1,${MASTER_ADDR}" -ray start --head --node-ip-address ${MASTER_ADDR} --num-gpus 8 --disable-usage-stats --dashboard-host=0.0.0.0 --dashboard-port=8265 -for WORKER_IP in $(awk '{print $1}' /root/mpi_rack_hostfile); do - if [[ "$WORKER_IP" == "$MLP_WORKER_0_HOST" ]]; then - continue - fi - echo "Starting Ray worker on ${WORKER_IP}" - ssh root@"${WORKER_IP}" \ - "pkill -9 -f "vllm serve" ; ray stop --force ; pkill -9 python ; ray start --address=${MASTER_ADDR}:6379 --num-gpus 8 --node-ip-address ${WORKER_IP} --disable-usage-stats --dashboard-host=0.0.0.0 --dashboard-port=8265" & -done -wait - - -# Build the runtime environment JSON with proper variable substitution -RUNTIME_ENV_JSON="{ - \"env_vars\": { - \"PYTHONPATH\": \"/root/Megatron-LM/\", - \"CUDA_DEVICE_MAX_CONNECTIONS\": \"1\", - \"NCCL_NVLS_ENABLE\": \"${HAS_NVLINK}\", - \"no_proxy\": \"${no_proxy}\", - \"MASTER_ADDR\": \"${MASTER_ADDR}\", - \"PYTORCH_CUDA_ALLOC_CONF\": \"expandable_segments:True\" - } -}" - -ray job submit --address="http://127.0.0.1:8265" \ - --runtime-env-json="${RUNTIME_ENV_JSON}" \ - -- python3 train_async.py \ - --actor-num-nodes 4 \ - --actor-num-gpus-per-node 8 \ - ${MODEL_ARGS[@]} \ - ${CKPT_ARGS[@]} \ - ${SFT_ARGS[@]} \ - ${OPTIMIZER_ARGS[@]} \ - ${WANDB_ARGS[@]} \ - ${PERF_ARGS[@]} \ - ${EVAL_ARGS[@]} \ - ${MISC_ARGS[@]} diff --git a/scripts/run-qwen3-235B-A22B.sh b/scripts/run-qwen3-235B-A22B.sh deleted file mode 100644 index f9e828af4..000000000 --- a/scripts/run-qwen3-235B-A22B.sh +++ /dev/null @@ -1,182 +0,0 @@ -#!/bin/bash - -# for rerun the task -pkill -9 -f "vllm serve" -sleep 3 -ray stop --force -pkill -9 ray -pkill -9 python -sleep 3 -pkill -9 ray -pkill -9 python - -set -ex - -# if base folder not set raise error -if [ -z "${BASE_FOLDER}" ]; then - echo "BASE_FOLDER is not set. Please set it to the base directory of your checkpoints." - exit 1 -fi - -if [ -z "${MASTER_ADDR}" ]; then - echo "MASTER_ADDR is not set. Please set it to the master node address." - exit 1 -fi - -# will prevent ray from buffering stdout/stderr -export PYTHONBUFFERED=16 - -NVLINK_COUNT=$(nvidia-smi topo -m 2>/dev/null | grep -o 'NV[0-9][0-9]*' | wc -l) -if [ "$NVLINK_COUNT" -gt 0 ]; then - HAS_NVLINK=1 -else - HAS_NVLINK=0 -fi -echo "HAS_NVLINK: $HAS_NVLINK (detected $NVLINK_COUNT NVLink references)" - -SCRIPT_DIR="$(cd -- "$(dirname -- "${BASH_SOURCE[0]}")" &>/dev/null && pwd)" -source "${SCRIPT_DIR}/models/qwen3-235B-A22B.sh" - -CKPT_ARGS=( - --hf-checkpoint ${BASE_FOLDER}/Qwen3-235B-A22B-FP8 - --ref-load ${BASE_FOLDER}/Qwen3-235B-A22B_torch_dist - --load ${BASE_FOLDER}/Qwen3-235B-A22B_slime/ - --save ${BASE_FOLDER}/Qwen3-235B-A22B_slime/ - --save-interval 20 -) - -ROLLOUT_ARGS=( - --prompt-data ${BASE_FOLDER}/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 8 - --n-samples-per-prompt 8 - --rollout-max-response-len 8192 - --rollout-temperature 1 - - --global-batch-size 64 - --balance-data -) - -EVAL_ARGS=( - #--eval-interval 20 - --eval-prompt-data aime ${BASE_FOLDER}/aime-2024/aime-2024.jsonl - --n-samples-per-eval-prompt 16 - --eval-max-response-len 16384 - --eval-top-p 1 -) - -PERF_ARGS=( - --tensor-model-parallel-size 4 - --sequence-parallel - --pipeline-model-parallel-size 4 - --context-parallel-size 2 - --expert-model-parallel-size 16 - --expert-tensor-parallel-size 1 - --decoder-last-pipeline-num-layers 22 - - --recompute-granularity full - --recompute-method uniform - --recompute-num-layers 1 - - # --micro-batch-size 1 - --use-dynamic-batch-size - --max-tokens-per-gpu 16384 -) - -GRPO_ARGS=( - --advantage-estimator gspo - #--use-kl-loss - --kl-loss-coef 0.00 - --kl-loss-type low_var_kl - --entropy-coef 0.00 - --eps-clip 4e-4 -) - -OPTIMIZER_ARGS=( - --optimizer adam - --lr 1e-6 - --lr-decay-style constant - --weight-decay 0.1 - --adam-beta1 0.9 - --adam-beta2 0.98 - - --optimizer-cpu-offload - --overlap-cpu-optimizer-d2h-h2d - --use-precision-aware-optimizer -) - -WANDB_ARGS=( - # --use-wandb - # --wandb-project vime-dev - # --wandb-group qwen3-235B-A22B -) - -VLLM_ARGS=( - --rollout-num-gpus-per-engine 32 - --vllm-gpu-memory-utilization 0.7 - --vllm-data-parallel-size 4 - --vllm-enable-expert-parallel - --vllm-cudagraph-capture-sizes 1 2 4 8 $(seq 16 8 256) - --vllm-all2all-backend deepep_high_throughput -) - -MISC_ARGS=( - # default dropout in megatron is 0.1 - --attention-dropout 0.0 - --hidden-dropout 0.0 - # should be good for model performance - --accumulate-allreduce-grads-in-fp32 - --attention-softmax-in-fp32 - # need to comment this when using model with MLA - --attention-backend flash -) - -# launch the master node of ray in container -export no_proxy="127.0.0.1,${MASTER_ADDR}" -ray start --head --node-ip-address ${MASTER_ADDR} --num-gpus 8 --disable-usage-stats --dashboard-host=0.0.0.0 --dashboard-port=8265 -for WORKER_IP in $(awk '{print $1}' /root/mpi_rack_hostfile); do - if [[ "$WORKER_IP" == "$MLP_WORKER_0_HOST" ]]; then - continue - fi - echo "Starting Ray worker on ${WORKER_IP}" - ssh root@"${WORKER_IP}" \ - "pkill -9 -f "vllm serve" ; ray stop --force ; pkill -9 python ; ray start --address=${MASTER_ADDR}:6379 --num-gpus 8 --node-ip-address ${WORKER_IP} --disable-usage-stats --dashboard-host=0.0.0.0 --dashboard-port=8265" & -done -wait - - -# Build the runtime environment JSON with proper variable substitution -RUNTIME_ENV_JSON="{ - \"env_vars\": { - \"PYTHONPATH\": \"/root/Megatron-LM/\", - \"CUDA_DEVICE_MAX_CONNECTIONS\": \"1\", - \"NCCL_NVLS_ENABLE\": \"${HAS_NVLINK}\", - \"no_proxy\": \"${no_proxy}\", - \"MASTER_ADDR\": \"${MASTER_ADDR}\" - } -}" - -ray job submit --address="http://127.0.0.1:8265" \ - --runtime-env-json="${RUNTIME_ENV_JSON}" \ - -- python3 train.py \ - --actor-num-nodes 8 \ - --actor-num-gpus-per-node 8 \ - --rollout-num-gpus 64 \ - --update-weight-buffer-size $(( 1024 * 1024 * 1024 * 4 )) \ - ${MODEL_ARGS[@]} \ - ${CKPT_ARGS[@]} \ - ${ROLLOUT_ARGS[@]} \ - ${OPTIMIZER_ARGS[@]} \ - ${GRPO_ARGS[@]} \ - ${WANDB_ARGS[@]} \ - ${PERF_ARGS[@]} \ - ${EVAL_ARGS[@]} \ - ${VLLM_ARGS[@]} \ - ${MISC_ARGS[@]} diff --git a/scripts/run-qwen3-32B.sh b/scripts/run-qwen3-32B.sh deleted file mode 100644 index 82afd1af4..000000000 --- a/scripts/run-qwen3-32B.sh +++ /dev/null @@ -1,154 +0,0 @@ -#!/bin/bash - -# for rerun the task -pkill -9 -f "vllm serve" -sleep 3 -ray stop --force -pkill -9 ray -pkill -9 python -sleep 3 -pkill -9 ray -pkill -9 python - -set -ex - -# will prevent ray from buffering stdout/stderr -export PYTHONBUFFERED=16 - -NVLINK_COUNT=$(nvidia-smi topo -m 2>/dev/null | grep -o 'NV[0-9][0-9]*' | wc -l) -if [ "$NVLINK_COUNT" -gt 0 ]; then - HAS_NVLINK=1 -else - HAS_NVLINK=0 -fi -echo "HAS_NVLINK: $HAS_NVLINK (detected $NVLINK_COUNT NVLink references)" - -SCRIPT_DIR="$(cd -- "$(dirname -- "${BASH_SOURCE[0]}")" &>/dev/null && pwd)" -source "${SCRIPT_DIR}/models/qwen3-32B.sh" - -CKPT_ARGS=( - --hf-checkpoint /root/Qwen3-32B - --ref-load /root/Qwen3-32B_torch_dist/ - --load /root/Qwen3-32B_slime - --save /root/Qwen3-32B_slime - --save-interval 20 -) - -ROLLOUT_ARGS=( - --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 5 - --rollout-batch-size 32 - --n-samples-per-prompt 8 - --rollout-max-response-len 8192 - --rollout-temperature 1 - - --global-batch-size 256 - --balance-data -) - -EVAL_ARGS=( - --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 -) - -PERF_ARGS=( - --tensor-model-parallel-size 8 - --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 - - # --micro-batch-size 1 - --use-dynamic-batch-size - --max-tokens-per-gpu 20480 -) - -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 -) - -OPTIMIZER_ARGS=( - --optimizer adam - --lr 1e-6 - --lr-decay-style constant - --weight-decay 0.1 - --adam-beta1 0.9 - --adam-beta2 0.98 - - --optimizer-cpu-offload - --overlap-cpu-optimizer-d2h-h2d - --use-precision-aware-optimizer -) - -WANDB_ARGS=( - #--use-wandb - # --wandb-project vime-dev - # --wandb-group qwen3-30B-A3B-test - # --wandb-key ${WANDB_KEY} -) - -VLLM_ARGS=( - --rollout-num-gpus-per-engine 8 - --vllm-gpu-memory-utilization 0.7 - --vllm-cudagraph-capture-sizes 1 2 4 8 $(seq 16 8 256) -) - -MISC_ARGS=( - # default dropout in megatron is 0.1 - --attention-dropout 0.0 - --hidden-dropout 0.0 - # should be good for model performance - --accumulate-allreduce-grads-in-fp32 - --attention-softmax-in-fp32 - # need to comment this when using model with MLA - --attention-backend flash -) - -# launch the master node of ray in container -export MASTER_ADDR=${MASTER_ADDR:-"127.0.0.1"} -ray start --head --node-ip-address ${MASTER_ADDR} --num-gpus 8 --disable-usage-stats --dashboard-host=0.0.0.0 --dashboard-port=8265 - -# Build the runtime environment JSON with proper variable substitution -RUNTIME_ENV_JSON="{ - \"env_vars\": { - \"PYTHONPATH\": \"/root/Megatron-LM/\", - \"CUDA_DEVICE_MAX_CONNECTIONS\": \"1\", - \"NCCL_NVLS_ENABLE\": \"${HAS_NVLINK}\" - } -}" - -ray job submit --address="http://127.0.0.1:8265" \ - --runtime-env-json="${RUNTIME_ENV_JSON}" \ - -- python3 train.py \ - --actor-num-nodes 1 \ - --actor-num-gpus-per-node 8 \ - --colocate \ - ${MODEL_ARGS[@]} \ - ${CKPT_ARGS[@]} \ - ${ROLLOUT_ARGS[@]} \ - ${OPTIMIZER_ARGS[@]} \ - ${GRPO_ARGS[@]} \ - ${WANDB_ARGS[@]} \ - ${PERF_ARGS[@]} \ - ${EVAL_ARGS[@]} \ - ${VLLM_ARGS[@]} \ - ${MISC_ARGS[@]} diff --git a/scripts/run-qwen3-4B-base-sft.sh b/scripts/run-qwen3-4B-base-sft.sh deleted file mode 100644 index 49d3326d7..000000000 --- a/scripts/run-qwen3-4B-base-sft.sh +++ /dev/null @@ -1,127 +0,0 @@ -#!/bin/bash - -# for rerun the task -pkill -9 -f "vllm serve" -sleep 3 -ray stop --force -pkill -9 ray -pkill -9 python -sleep 3 -pkill -9 ray -pkill -9 python - -set -ex - -# will prevent ray from buffering stdout/stderr -export PYTHONBUFFERED=16 - -NVLINK_COUNT=$(nvidia-smi topo -m 2>/dev/null | grep -o 'NV[0-9][0-9]*' | wc -l) -if [ "$NVLINK_COUNT" -gt 0 ]; then - HAS_NVLINK=1 -else - HAS_NVLINK=0 -fi -echo "HAS_NVLINK: $HAS_NVLINK (detected $NVLINK_COUNT NVLink references)" - -SCRIPT_DIR="$(cd -- "$(dirname -- "${BASH_SOURCE[0]}")" &>/dev/null && pwd)" -source "${SCRIPT_DIR}/models/qwen3-4B.sh" - -CKPT_ARGS=( - --hf-checkpoint /root/Qwen3-4B-Base/ - --ref-load /root/Qwen3-4B-Base_torch_dist - --load /root/Qwen3-4B-Base_slime/ - --save /root/Qwen3-4B-Base_slime/ - --save-interval 1000 -) - -SFT_ARGS=( - --rollout-function-path vime.rollout.sft_rollout.generate_rollout - --prompt-data /root/openhermes2_5.parquet - --input-key messages - # --apply-chat-template - --rollout-shuffle - --num-epoch 3 - --rollout-batch-size 128 - --global-batch-size 128 - - --loss-type sft_loss - --calculate-per-token-loss - --disable-compute-advantages-and-returns - --debug-train-only -) - -PERF_ARGS=( - --tensor-model-parallel-size 1 - --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 - - # --micro-batch-size 1 - --use-dynamic-batch-size - --max-tokens-per-gpu 9216 -) - -OPTIMIZER_ARGS=( - --optimizer adam - --lr 1e-5 - --lr-decay-style cosine - --min-lr 1e-6 - --lr-warmup-fraction 0.1 - --weight-decay 0.1 - --adam-beta1 0.9 - --adam-beta2 0.95 -) - -WANDB_ARGS=( - # --use-wandb - # --wandb-project vime-dev - # --wandb-group qwen3-4B-base-sft - # --wandb-key ${WANDB_KEY} -) - -MISC_ARGS=( - # default dropout in megatron is 0.1 - --attention-dropout 0.0 - --hidden-dropout 0.0 - # should be good for model performance - --accumulate-allreduce-grads-in-fp32 - --attention-softmax-in-fp32 - # need to comment this when using model with MLA - --attention-backend flash -) - -# launch the master node of ray in container -export MASTER_ADDR=${MASTER_ADDR:-"127.0.0.1"} -export no_proxy="127.0.0.1,${MASTER_ADDR}" -ray start --head --node-ip-address ${MASTER_ADDR} --num-gpus 8 --disable-usage-stats --dashboard-host=0.0.0.0 --dashboard-port=8265 - - -# Build the runtime environment JSON with proper variable substitution -RUNTIME_ENV_JSON="{ - \"env_vars\": { - \"PYTHONPATH\": \"/root/Megatron-LM/\", - \"CUDA_DEVICE_MAX_CONNECTIONS\": \"1\", - \"NCCL_NVLS_ENABLE\": \"${HAS_NVLINK}\", - \"PYTORCH_CUDA_ALLOC_CONF\": \"expandable_segments:True\" - } -}" - -ray job submit --address="http://127.0.0.1:8265" \ - --runtime-env-json="${RUNTIME_ENV_JSON}" \ - -- python3 train_async.py \ - --actor-num-nodes 1 \ - --actor-num-gpus-per-node 8 \ - ${MODEL_ARGS[@]} \ - ${CKPT_ARGS[@]} \ - ${SFT_ARGS[@]} \ - ${OPTIMIZER_ARGS[@]} \ - ${WANDB_ARGS[@]} \ - ${PERF_ARGS[@]} \ - ${EVAL_ARGS[@]} \ - ${MISC_ARGS[@]} diff --git a/scripts/run-qwen3-next-80B-A3B.sh b/scripts/run-qwen3-next-80B-A3B.sh deleted file mode 100644 index 3e942e091..000000000 --- a/scripts/run-qwen3-next-80B-A3B.sh +++ /dev/null @@ -1,194 +0,0 @@ -#!/bin/bash - -# for rerun the task -pkill -9 -f "vllm serve" -sleep 3 -ray stop --force -pkill -9 ray -pkill -9 python -sleep 3 -pkill -9 ray -pkill -9 python - -set -ex - -if [ -z "${BASE_FOLDER:-}" ]; then - echo "BASE_FOLDER is not set. Please set it to the base directory of your checkpoints." - exit 1 -fi - -MASTER_ADDR=${MASTER_ADDR:-} -if [ -z "${MASTER_ADDR}" ]; then - echo "MASTER_ADDR is not set. Please set it to the master node address." - exit 1 -fi - -# will prevent ray from buffering stdout/stderr -export PYTHONBUFFERED=16 - -# unset proxy to avoid distributed startup issues -unset http_proxy https_proxy HTTP_PROXY HTTPS_PROXY - -ACTOR_NUM_NODES=${ACTOR_NUM_NODES:-4} -ACTOR_NUM_GPUS_PER_NODE=${ACTOR_NUM_GPUS_PER_NODE:-8} -CP_SIZE=${CP_SIZE:-4} -SOCKET_IFNAME=${SOCKET_IFNAME:-eth0} - -NVLINK_COUNT=$(nvidia-smi topo -m 2>/dev/null | grep -o 'NV[0-9][0-9]*' | wc -l) -if [ "$NVLINK_COUNT" -gt 0 ]; then - HAS_NVLINK=1 -else - HAS_NVLINK=0 -fi -echo "HAS_NVLINK: $HAS_NVLINK (detected $NVLINK_COUNT NVLink references)" - -SCRIPT_DIR="$(cd -- "$(dirname -- "${BASH_SOURCE[0]}")" &>/dev/null && pwd)" -source "${SCRIPT_DIR}/models/qwen3-next-80B-A3B.sh" - -CKPT_ARGS=( - --hf-checkpoint "${BASE_FOLDER}/Qwen3-Next-80B-A3B-Thinking" - --ref-load "${BASE_FOLDER}/Qwen3-Next-80B-A3B-Thinking_torch_dist" - --load "${BASE_FOLDER}/Qwen3-Next-80B-A3B-Thinking_slime/" - --save "${BASE_FOLDER}/Qwen3-Next-80B-A3B-Thinking_slime/" - --save-interval 20 -) - -ROLLOUT_ARGS=( - --prompt-data "${BASE_FOLDER}/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 8 - --n-samples-per-prompt 8 - --rollout-max-response-len 32768 - --rollout-temperature 1 - --global-batch-size 64 - --balance-data -) - -EVAL_ARGS=( - --eval-interval 20 - --eval-prompt-data aime "${BASE_FOLDER}/aime-2024/aime-2024.jsonl" - --n-samples-per-eval-prompt 8 - --eval-max-response-len 32768 - --eval-top-p 1 -) - -PERF_ARGS=( - --tensor-model-parallel-size 2 - --sequence-parallel - --pipeline-model-parallel-size 4 - --decoder-last-pipeline-num-layers 9 - --context-parallel-size "${CP_SIZE}" - --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 8192 -) - -GRPO_ARGS=( - --advantage-estimator gspo - --kl-loss-coef 0.00 - --kl-loss-type low_var_kl - --kl-coef 0.00 - --entropy-coef 0.00 - --eps-clip 4e-4 -) - -OPTIMIZER_ARGS=( - --optimizer adam - --lr 1e-6 - --lr-decay-style constant - --weight-decay 0.1 - --adam-beta1 0.9 - --adam-beta2 0.98 - - --optimizer-cpu-offload - --overlap-cpu-optimizer-d2h-h2d - --use-precision-aware-optimizer -) - -WANDB_ARGS=( - # --use-wandb - # --wandb-project vime-dev - # --wandb-group qwen3-next-80B-A3B-32k - # --wandb-key ${WANDB_KEY} -) - -VLLM_ARGS=( - --rollout-num-gpus-per-engine 8 - --vllm-gpu-memory-utilization 0.8 - - # mtp - - --vllm-max-num-seqs 256 - --vllm-enable-expert-parallel - --vllm-cudagraph-capture-sizes 1 2 4 8 $(seq 16 8 128) - --vllm-speculative-config '{"method":"mtp","num_speculative_tokens":3}' -) - -MISC_ARGS=( - --attention-dropout 0.0 - --hidden-dropout 0.0 - --accumulate-allreduce-grads-in-fp32 - --attention-softmax-in-fp32 - --attention-backend flash - --moe-token-dispatcher-type flex - --moe-enable-deepep -) - -export no_proxy="127.0.0.1,${MASTER_ADDR}" -ray start --head --node-ip-address "${MASTER_ADDR}" --num-gpus "${ACTOR_NUM_GPUS_PER_NODE}" --disable-usage-stats --dashboard-host=0.0.0.0 --dashboard-port=8265 - -HOSTFILE=${HOSTFILE:-} -if [ -n "${HOSTFILE}" ]; then - for WORKER_IP in $(awk '{print $1}' "${HOSTFILE}"); do - if [[ "${WORKER_IP}" == "${MASTER_ADDR}" ]]; then - continue - fi - echo "Starting Ray worker on ${WORKER_IP}" - ssh root@"${WORKER_IP}" \ - "pkill -9 -f "vllm serve" ; ray stop --force ; pkill -9 python ; ray start --address=${MASTER_ADDR}:6379 --num-gpus ${ACTOR_NUM_GPUS_PER_NODE} --node-ip-address ${WORKER_IP} --disable-usage-stats" & - done - wait -fi - -RUNTIME_ENV_JSON=$(cat </dev/null | grep -o 'NV[0-9][0-9]*' | wc -l) -if [ "$NVLINK_COUNT" -gt 0 ]; then - HAS_NVLINK=1 -else - HAS_NVLINK=0 -fi -echo "HAS_NVLINK: $HAS_NVLINK (detected $NVLINK_COUNT NVLink references)" - -SCRIPT_DIR="$(cd -- "$(dirname -- "${BASH_SOURCE[0]}")" &>/dev/null && pwd)" -source "${SCRIPT_DIR}/models/qwen3.5-27B.sh" - -CKPT_ARGS=( - --hf-checkpoint "${BASE_FOLDER}/Qwen3.5-27B" - --ref-load "${BASE_FOLDER}/Qwen3.5-27B_torch_dist/" - --load "${BASE_FOLDER}/Qwen3.5-27B_slime/" - --save "${BASE_FOLDER}/Qwen3.5-27B_slime/" - --save-interval 20 -) - -ROLLOUT_ARGS=( - --prompt-data "${BASE_FOLDER}/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 8 - --n-samples-per-prompt 8 - --rollout-max-response-len 32768 - --rollout-temperature 1.0 - --global-batch-size 64 - --balance-data -) - -EVAL_ARGS=( - --eval-interval 20 - --eval-prompt-data aime "${BASE_FOLDER}/aime-2024/aime-2024.jsonl" - --n-samples-per-eval-prompt 8 - --eval-max-response-len 32768 - --eval-top-p 1 -) - -PERF_ARGS=( - --tensor-model-parallel-size 4 - --sequence-parallel - --pipeline-model-parallel-size 2 - --decoder-last-pipeline-num-layers 30 - --context-parallel-size "${CP_SIZE}" - --expert-model-parallel-size 1 - --expert-tensor-parallel-size 1 - - --recompute-granularity full - --recompute-method uniform - --recompute-num-layers 1 - - --use-dynamic-batch-size - --calculate-per-token-loss - --max-tokens-per-gpu 8192 -) - -GRPO_ARGS=( - --advantage-estimator grpo - --kl-loss-coef 0.00 - --kl-loss-type low_var_kl - --kl-coef 0.00 - --entropy-coef 0.00 - --eps-clip 0.2 -) - -OPTIMIZER_ARGS=( - --optimizer adam - --lr 1e-6 - --lr-decay-style constant - --weight-decay 0.1 - --adam-beta1 0.9 - --adam-beta2 0.98 - --optimizer-cpu-offload - --overlap-cpu-optimizer-d2h-h2d - --use-precision-aware-optimizer -) - -WANDB_ARGS=( - # --use-wandb - # --wandb-project vime-dev - # --wandb-group qwen3.5-27B-32k - # --wandb-key ${WANDB_KEY} -) - -VLLM_ARGS=( - --rollout-num-gpus-per-engine 2 - --vllm-gpu-memory-utilization 0.75 - --vllm-speculative-config '{"method":"mtp","num_speculative_tokens":3}' -) - -MISC_ARGS=( - --attention-dropout 0.0 - --hidden-dropout 0.0 - --accumulate-allreduce-grads-in-fp32 - --attention-softmax-in-fp32 - --attention-backend flash -) - -export no_proxy="127.0.0.1,${MASTER_ADDR}" -ray start --head --node-ip-address "${MASTER_ADDR}" --num-gpus "${ACTOR_NUM_GPUS_PER_NODE}" --disable-usage-stats --dashboard-host=0.0.0.0 --dashboard-port=8265 - -HOSTFILE=${HOSTFILE:-} -if [ -n "${HOSTFILE}" ]; then - for WORKER_IP in $(awk '{print $1}' "${HOSTFILE}"); do - if [[ "${WORKER_IP}" == "${MASTER_ADDR}" ]]; then - continue - fi - echo "Starting Ray worker on ${WORKER_IP}" - ssh root@"${WORKER_IP}" \ - "pkill -9 -f "vllm serve" ; ray stop --force ; pkill -9 python ; ray start --address=${MASTER_ADDR}:6379 --num-gpus ${ACTOR_NUM_GPUS_PER_NODE} --node-ip-address ${WORKER_IP} --disable-usage-stats" & - done - wait -fi - -RUNTIME_ENV_JSON=$(cat </dev/null | grep -o 'NV[0-9][0-9]*' | wc -l) -if [ "$NVLINK_COUNT" -gt 0 ]; then - HAS_NVLINK=1 -else - HAS_NVLINK=0 -fi -echo "HAS_NVLINK: $HAS_NVLINK (detected $NVLINK_COUNT NVLink references)" - -SCRIPT_DIR="$(cd -- "$(dirname -- "${BASH_SOURCE[0]}")" &>/dev/null && pwd)" -source "${SCRIPT_DIR}/models/qwen3.5-35B-A3B.sh" - -CKPT_ARGS=( - --hf-checkpoint ${BASE_FOLDER}/Qwen3.5-35B-A3B - --ref-load ${BASE_FOLDER}/Qwen3.5-35B-A3B_torch_dist - --load ${BASE_FOLDER}/Qwen3.5-35B-A3B_slime/ - --save ${BASE_FOLDER}/Qwen3.5-35B-A3B_slime/ - --save-interval 20 -) - -SFT_ARGS=( - --rollout-function-path vime.rollout.sft_rollout.generate_rollout - --prompt-data ${BASE_FOLDER}/openhermes2_5.parquet - --input-key messages - --rollout-shuffle - --num-epoch 3 - --rollout-batch-size 128 - --global-batch-size 128 - - --loss-type sft_loss - --loss-mask-type qwen3_5 - --calculate-per-token-loss - --disable-compute-advantages-and-returns - --debug-train-only -) - -PERF_ARGS=( - --tensor-model-parallel-size 2 - --sequence-parallel - --pipeline-model-parallel-size 1 - --context-parallel-size 1 - --expert-model-parallel-size 8 - --expert-tensor-parallel-size 1 - - --recompute-granularity full - --recompute-method uniform - --recompute-num-layers 1 - - # --micro-batch-size 1 - --use-dynamic-batch-size - --max-tokens-per-gpu 8192 -) - -OPTIMIZER_ARGS=( - --optimizer adam - --lr 1e-5 - --lr-decay-style cosine - --min-lr 1e-6 - --lr-warmup-fraction 0.1 - --weight-decay 0.1 - --adam-beta1 0.9 - --adam-beta2 0.98 - - --use-distributed-optimizer - --optimizer-cpu-offload - --overlap-cpu-optimizer-d2h-h2d - --use-precision-aware-optimizer -) - -WANDB_ARGS=( - # --use-wandb - # --wandb-project vime-dev - # --wandb-group qwen3.5-35B-sft -) - -MISC_ARGS=( - # default dropout in megatron is 0.1 - --attention-dropout 0.0 - --hidden-dropout 0.0 - # should be good for model performance - --accumulate-allreduce-grads-in-fp32 - --attention-softmax-in-fp32 - # need to comment this when using model with MLA - --attention-backend flash - - --moe-token-dispatcher-type flex - --moe-enable-deepep -) - -SPEC_ARGS=( -# --mtp-num-layers 1 -# --enable-mtp-training -# --mtp-loss-scaling-factor 0.1 -) - -# launch the master node of ray in container -export no_proxy="127.0.0.1,${MASTER_ADDR}" -ray start --head --node-ip-address ${MASTER_ADDR} --num-gpus 8 --disable-usage-stats --dashboard-host=0.0.0.0 --dashboard-port=8265 -for WORKER_IP in $(awk '{print $1}' /root/mpi_rack_hostfile); do - if [[ "$WORKER_IP" == "$MLP_WORKER_0_HOST" ]]; then - continue - fi - echo "Starting Ray worker on ${WORKER_IP}" - ssh root@"${WORKER_IP}" \ - "pkill -9 -f "vllm serve" ; ray stop --force ; pkill -9 python ; ray start --address=${MASTER_ADDR}:6379 --num-gpus 8 --node-ip-address ${WORKER_IP} --disable-usage-stats --dashboard-host=0.0.0.0 --dashboard-port=8265" & -done -wait - - -# Build the runtime environment JSON with proper variable substitution -RUNTIME_ENV_JSON="{ - \"env_vars\": { - \"PYTHONPATH\": \"/root/Megatron-LM/\", - \"CUDA_DEVICE_MAX_CONNECTIONS\": \"1\", - \"NCCL_NVLS_ENABLE\": \"${HAS_NVLINK}\", - \"no_proxy\": \"${no_proxy}\", - \"MASTER_ADDR\": \"${MASTER_ADDR}\", - \"PYTORCH_CUDA_ALLOC_CONF\": \"expandable_segments:True\" - } -}" - -ray job submit --address="http://127.0.0.1:8265" \ - --runtime-env-json="${RUNTIME_ENV_JSON}" \ - -- python3 train_async.py \ - --actor-num-nodes 1 \ - --actor-num-gpus-per-node 8 \ - ${MODEL_ARGS[@]} \ - ${CKPT_ARGS[@]} \ - ${SFT_ARGS[@]} \ - ${OPTIMIZER_ARGS[@]} \ - ${WANDB_ARGS[@]} \ - ${PERF_ARGS[@]} \ - ${EVAL_ARGS[@]} \ - ${MISC_ARGS[@]} \ - ${SPEC_ARGS[@]}