初始化项目,由ModelHub XC社区提供模型
Model: ayh015/myLightningOPD Source: Original Platform
This commit is contained in:
145
configs/lightning_opd/qwen3-30b-a3b-lightning-opd.py
Normal file
145
configs/lightning_opd/qwen3-30b-a3b-lightning-opd.py
Normal file
@@ -0,0 +1,145 @@
|
||||
# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
|
||||
# SPDX-License-Identifier: Apache-2.0
|
||||
|
||||
import os
|
||||
from pathlib import Path
|
||||
|
||||
import slime.utils.external_utils.command_utils as U
|
||||
|
||||
# Lightning OPD: response tokens + teacher logprobs are pre-computed in parquet.
|
||||
# No teacher server needed at training time -> all 8 GPUs go to the actor.
|
||||
#
|
||||
# Required env vars:
|
||||
# SFT_CHECKPOINT - path to the SFT checkpoint (HF format)
|
||||
# LIGHTNING_OPD_DATA - path to the precomputed parquet
|
||||
|
||||
MODEL_NAME = "Qwen3-30B-A3B-Base-Open-Thoughts-Qwen3-30B-A3B-Thinking-2507-sft"
|
||||
MODEL_TYPE = "qwen3-30B-A3B"
|
||||
NUM_GPUS = 8
|
||||
SFT_CHECKPOINT = os.environ["SFT_CHECKPOINT"]
|
||||
|
||||
|
||||
def prepare():
|
||||
U.convert_checkpoint(
|
||||
model_name=MODEL_NAME,
|
||||
megatron_model_type=MODEL_TYPE,
|
||||
num_gpus_per_node=NUM_GPUS,
|
||||
hf_checkpoint=SFT_CHECKPOINT,
|
||||
)
|
||||
|
||||
|
||||
def execute(rerun=True):
|
||||
load_save_path = f"/root/models/{MODEL_NAME}_ckpt__{Path(__file__).stem}/"
|
||||
|
||||
ckpt_args = (
|
||||
f"--hf-checkpoint {SFT_CHECKPOINT} "
|
||||
f"--ref-load /root/models/{MODEL_NAME}_torch_dist "
|
||||
f"--load {load_save_path} "
|
||||
f"--save {load_save_path} "
|
||||
"--save-interval 10 "
|
||||
"--save-retain-interval 100 "
|
||||
)
|
||||
|
||||
rollout_args = (
|
||||
f"--prompt-data {os.environ['LIGHTNING_OPD_DATA']} "
|
||||
"--input-key prompt "
|
||||
"--label-key label "
|
||||
"--rollout-shuffle "
|
||||
"--num-rollout 3000 "
|
||||
"--rollout-batch-size 256 "
|
||||
"--n-samples-per-prompt 1 "
|
||||
"--rollout-max-response-len 4096 "
|
||||
"--global-batch-size 256 "
|
||||
)
|
||||
|
||||
rm_args = (
|
||||
"--custom-rm-path slime.rollout.on_policy_distillation.reward_func "
|
||||
"--custom-reward-post-process-path slime.rollout.on_policy_distillation.post_process_rewards "
|
||||
"--include-verifiable-reward "
|
||||
)
|
||||
|
||||
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 "
|
||||
"--use-dynamic-batch-size "
|
||||
"--max-tokens-per-gpu 12288 "
|
||||
)
|
||||
|
||||
grpo_args = (
|
||||
"--advantage-estimator on_policy_distillation "
|
||||
"--use-kl-loss "
|
||||
"--kl-loss-coef 0.00 "
|
||||
"--kl-loss-type low_var_kl "
|
||||
"--entropy-coef 0.00 "
|
||||
)
|
||||
|
||||
optimizer_args = (
|
||||
"--optimizer adam "
|
||||
"--lr 2e-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 = ""
|
||||
if os.environ.get("WANDB_KEY"):
|
||||
wandb_args = (
|
||||
"--use-wandb "
|
||||
"--wandb-project lightning-opd "
|
||||
f"--wandb-group {Path(__file__).stem} "
|
||||
f"--wandb-key {os.environ['WANDB_KEY']} "
|
||||
)
|
||||
|
||||
sglang_args = (
|
||||
"--rollout-num-gpus-per-engine 8 "
|
||||
"--sglang-mem-fraction-static 0.8 "
|
||||
"--sglang-cuda-graph-max-bs 32 "
|
||||
"--sglang-max-running-requests 512 "
|
||||
"--sglang-enable-metrics "
|
||||
)
|
||||
|
||||
misc_args = (
|
||||
"--attention-dropout 0.0 "
|
||||
"--hidden-dropout 0.0 "
|
||||
"--accumulate-allreduce-grads-in-fp32 "
|
||||
"--attention-softmax-in-fp32 "
|
||||
"--attention-backend flash "
|
||||
"--actor-num-nodes 1 "
|
||||
"--actor-num-gpus-per-node 8 "
|
||||
"--rollout-num-gpus 0 "
|
||||
)
|
||||
|
||||
train_args = (
|
||||
f"{ckpt_args} "
|
||||
f"{rollout_args} "
|
||||
f"{rm_args} "
|
||||
f"{grpo_args} "
|
||||
f"{optimizer_args} "
|
||||
f"{wandb_args} "
|
||||
f"{perf_args} "
|
||||
f"{sglang_args} "
|
||||
f"{misc_args} "
|
||||
)
|
||||
|
||||
U.execute_train(
|
||||
rerun=rerun,
|
||||
train_args=train_args,
|
||||
num_gpus_per_node=NUM_GPUS,
|
||||
megatron_model_type=MODEL_TYPE,
|
||||
)
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
prepare()
|
||||
execute(rerun=False)
|
||||
140
configs/lightning_opd/qwen3-4b-lightning-opd.py
Normal file
140
configs/lightning_opd/qwen3-4b-lightning-opd.py
Normal file
@@ -0,0 +1,140 @@
|
||||
# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
|
||||
# SPDX-License-Identifier: Apache-2.0
|
||||
|
||||
import os
|
||||
from pathlib import Path
|
||||
|
||||
import slime.utils.external_utils.command_utils as U
|
||||
|
||||
# Lightning OPD: response tokens + teacher logprobs are pre-computed in parquet.
|
||||
# No teacher server needed at training time -> all 8 GPUs go to the actor.
|
||||
#
|
||||
# Required env vars:
|
||||
# SFT_CHECKPOINT - path to the SFT checkpoint (HF format)
|
||||
# LIGHTNING_OPD_DATA - path to the precomputed parquet
|
||||
|
||||
MODEL_NAME = "Qwen3-4B-Base-Open-Thoughts-Qwen3-8B-sft-3k"
|
||||
MODEL_TYPE = "qwen3-4B"
|
||||
NUM_GPUS = 4
|
||||
SFT_CHECKPOINT = os.environ["SFT_CHECKPOINT"]
|
||||
|
||||
|
||||
def prepare():
|
||||
U.convert_checkpoint(
|
||||
model_name=MODEL_NAME,
|
||||
megatron_model_type=MODEL_TYPE,
|
||||
num_gpus_per_node=NUM_GPUS,
|
||||
hf_checkpoint=SFT_CHECKPOINT,
|
||||
)
|
||||
|
||||
|
||||
def execute(rerun=True):
|
||||
load_save_path = f"/root/models/{MODEL_NAME}_ckpt__{Path(__file__).stem}/"
|
||||
|
||||
ckpt_args = (
|
||||
f"--hf-checkpoint {SFT_CHECKPOINT} "
|
||||
f"--ref-load /root/models/{MODEL_NAME}_torch_dist "
|
||||
f"--load {load_save_path} "
|
||||
f"--save {load_save_path} "
|
||||
"--save-interval 10 "
|
||||
"--save-retain-interval 100 "
|
||||
)
|
||||
|
||||
rollout_args = (
|
||||
f"--prompt-data {os.environ['LIGHTNING_OPD_DATA']} "
|
||||
"--input-key prompt "
|
||||
"--label-key label "
|
||||
"--rollout-shuffle "
|
||||
"--num-rollout 150 "
|
||||
"--rollout-batch-size 256 "
|
||||
"--n-samples-per-prompt 1 "
|
||||
"--rollout-max-response-len 4096 "
|
||||
"--global-batch-size 256 "
|
||||
"--rollout-temperature 0.8"
|
||||
)
|
||||
|
||||
rm_args = (
|
||||
"--custom-rm-path slime.rollout.on_policy_distillation.reward_func "
|
||||
"--custom-reward-post-process-path slime.rollout.on_policy_distillation.post_process_rewards "
|
||||
"--include-verifiable-reward "
|
||||
)
|
||||
|
||||
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 "
|
||||
"--use-dynamic-batch-size "
|
||||
"--max-tokens-per-gpu 16384 "
|
||||
)
|
||||
|
||||
grpo_args = (
|
||||
"--advantage-estimator on_policy_distillation "
|
||||
"--use-kl-loss "
|
||||
"--kl-loss-coef 0.00 "
|
||||
"--kl-loss-type low_var_kl "
|
||||
"--entropy-coef 0.00 "
|
||||
)
|
||||
|
||||
optimizer_args = (
|
||||
"--optimizer adam "
|
||||
"--lr 2e-6 "
|
||||
"--lr-decay-style constant "
|
||||
"--weight-decay 0.1 "
|
||||
"--adam-beta1 0.9 "
|
||||
"--adam-beta2 0.98 "
|
||||
)
|
||||
|
||||
wandb_args = ""
|
||||
if os.environ.get("WANDB_KEY"):
|
||||
wandb_args = (
|
||||
"--use-wandb "
|
||||
"--wandb-project lightning-opd "
|
||||
f"--wandb-group {Path(__file__).stem} "
|
||||
f"--wandb-key {os.environ['WANDB_KEY']} "
|
||||
)
|
||||
|
||||
sglang_args = (
|
||||
"--rollout-num-gpus-per-engine 1 "
|
||||
"--sglang-mem-fraction-static 0.4 "
|
||||
)
|
||||
|
||||
misc_args = (
|
||||
"--attention-dropout 0.0 "
|
||||
"--hidden-dropout 0.0 "
|
||||
"--accumulate-allreduce-grads-in-fp32 "
|
||||
"--attention-softmax-in-fp32 "
|
||||
"--attention-backend flash "
|
||||
"--actor-num-nodes 1 "
|
||||
"--actor-num-gpus-per-node 4 "
|
||||
"--rollout-num-gpus 0 "
|
||||
)
|
||||
|
||||
train_args = (
|
||||
f"{ckpt_args} "
|
||||
f"{rollout_args} "
|
||||
f"{rm_args} "
|
||||
f"{grpo_args} "
|
||||
f"{optimizer_args} "
|
||||
f"{wandb_args} "
|
||||
f"{perf_args} "
|
||||
f"{sglang_args} "
|
||||
f"{misc_args} "
|
||||
)
|
||||
|
||||
U.execute_train(
|
||||
rerun=rerun,
|
||||
train_args=train_args,
|
||||
num_gpus_per_node=NUM_GPUS,
|
||||
megatron_model_type=MODEL_TYPE,
|
||||
)
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
prepare()
|
||||
execute(rerun=False)
|
||||
139
configs/lightning_opd/qwen3-8b-lightning-opd.py
Normal file
139
configs/lightning_opd/qwen3-8b-lightning-opd.py
Normal file
@@ -0,0 +1,139 @@
|
||||
# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
|
||||
# SPDX-License-Identifier: Apache-2.0
|
||||
|
||||
import os
|
||||
from pathlib import Path
|
||||
|
||||
import slime.utils.external_utils.command_utils as U
|
||||
|
||||
# Lightning OPD: response tokens + teacher logprobs are pre-computed in parquet.
|
||||
# No teacher server needed at training time -> all 8 GPUs go to the actor.
|
||||
#
|
||||
# Required env vars:
|
||||
# SFT_CHECKPOINT - path to the SFT checkpoint (HF format)
|
||||
# LIGHTNING_OPD_DATA - path to the precomputed parquet
|
||||
|
||||
MODEL_NAME = "Qwen3-8B-Base-Open-Thoughts-Qwen3-32B-sft-3k"
|
||||
MODEL_TYPE = "qwen3-8B"
|
||||
NUM_GPUS = 8
|
||||
SFT_CHECKPOINT = os.environ["SFT_CHECKPOINT"]
|
||||
|
||||
|
||||
def prepare():
|
||||
U.convert_checkpoint(
|
||||
model_name=MODEL_NAME,
|
||||
megatron_model_type=MODEL_TYPE,
|
||||
num_gpus_per_node=NUM_GPUS,
|
||||
hf_checkpoint=SFT_CHECKPOINT,
|
||||
)
|
||||
|
||||
|
||||
def execute(rerun=True):
|
||||
load_save_path = f"/root/models/{MODEL_NAME}_ckpt__{Path(__file__).stem}/"
|
||||
|
||||
ckpt_args = (
|
||||
f"--hf-checkpoint {SFT_CHECKPOINT} "
|
||||
f"--ref-load /root/models/{MODEL_NAME}_torch_dist "
|
||||
f"--load {load_save_path} "
|
||||
f"--save {load_save_path} "
|
||||
"--save-interval 10 "
|
||||
"--save-retain-interval 100 "
|
||||
)
|
||||
|
||||
rollout_args = (
|
||||
f"--prompt-data {os.environ['LIGHTNING_OPD_DATA']} "
|
||||
"--input-key prompt "
|
||||
"--label-key label "
|
||||
"--rollout-shuffle "
|
||||
"--num-rollout 3000 "
|
||||
"--rollout-batch-size 256 "
|
||||
"--n-samples-per-prompt 1 "
|
||||
"--rollout-max-response-len 4096 "
|
||||
"--global-batch-size 256 "
|
||||
)
|
||||
|
||||
rm_args = (
|
||||
"--custom-rm-path slime.rollout.on_policy_distillation.reward_func "
|
||||
"--custom-reward-post-process-path slime.rollout.on_policy_distillation.post_process_rewards "
|
||||
"--include-verifiable-reward "
|
||||
)
|
||||
|
||||
perf_args = (
|
||||
"--tensor-model-parallel-size 4 "
|
||||
"--sequence-parallel "
|
||||
"--pipeline-model-parallel-size 1 "
|
||||
"--context-parallel-size 1 "
|
||||
"--expert-model-parallel-size 1 "
|
||||
"--expert-tensor-parallel-size 1 "
|
||||
"--recompute-granularity full "
|
||||
"--recompute-method uniform "
|
||||
"--recompute-num-layers 1 "
|
||||
"--use-dynamic-batch-size "
|
||||
"--max-tokens-per-gpu 16384 "
|
||||
)
|
||||
|
||||
grpo_args = (
|
||||
"--advantage-estimator on_policy_distillation "
|
||||
"--use-kl-loss "
|
||||
"--kl-loss-coef 0.00 "
|
||||
"--kl-loss-type low_var_kl "
|
||||
"--entropy-coef 0.00 "
|
||||
)
|
||||
|
||||
optimizer_args = (
|
||||
"--optimizer adam "
|
||||
"--lr 2e-6 "
|
||||
"--lr-decay-style constant "
|
||||
"--weight-decay 0.1 "
|
||||
"--adam-beta1 0.9 "
|
||||
"--adam-beta2 0.98 "
|
||||
)
|
||||
|
||||
wandb_args = ""
|
||||
if os.environ.get("WANDB_KEY"):
|
||||
wandb_args = (
|
||||
"--use-wandb "
|
||||
"--wandb-project lightning-opd "
|
||||
f"--wandb-group {Path(__file__).stem} "
|
||||
f"--wandb-key {os.environ['WANDB_KEY']} "
|
||||
)
|
||||
|
||||
sglang_args = (
|
||||
"--rollout-num-gpus-per-engine 1 "
|
||||
"--sglang-mem-fraction-static 0.4 "
|
||||
)
|
||||
|
||||
misc_args = (
|
||||
"--attention-dropout 0.0 "
|
||||
"--hidden-dropout 0.0 "
|
||||
"--accumulate-allreduce-grads-in-fp32 "
|
||||
"--attention-softmax-in-fp32 "
|
||||
"--attention-backend flash "
|
||||
"--actor-num-nodes 1 "
|
||||
"--actor-num-gpus-per-node 8 "
|
||||
"--rollout-num-gpus 0 "
|
||||
)
|
||||
|
||||
train_args = (
|
||||
f"{ckpt_args} "
|
||||
f"{rollout_args} "
|
||||
f"{rm_args} "
|
||||
f"{grpo_args} "
|
||||
f"{optimizer_args} "
|
||||
f"{wandb_args} "
|
||||
f"{perf_args} "
|
||||
f"{sglang_args} "
|
||||
f"{misc_args} "
|
||||
)
|
||||
|
||||
U.execute_train(
|
||||
rerun=rerun,
|
||||
train_args=train_args,
|
||||
num_gpus_per_node=NUM_GPUS,
|
||||
megatron_model_type=MODEL_TYPE,
|
||||
)
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
prepare()
|
||||
execute(rerun=False)
|
||||
49
configs/models/qwen3-30B-A3B.sh
Normal file
49
configs/models/qwen3-30B-A3B.sh
Normal file
@@ -0,0 +1,49 @@
|
||||
NLAYERS=48
|
||||
FIRST_K_DENSE_REPLACE=0
|
||||
|
||||
arr=()
|
||||
for ((i=0; i<NLAYERS; i++)); do
|
||||
if (( i < FIRST_K_DENSE_REPLACE )); then
|
||||
arr+=(0)
|
||||
else
|
||||
arr+=(1)
|
||||
fi
|
||||
done
|
||||
|
||||
printf -v MOE_LAYER_FREQ "[%s]" "$(IFS=', '; echo "${arr[*]}")"
|
||||
|
||||
|
||||
MODEL_ARGS=(
|
||||
--disable-bias-linear
|
||||
--qk-layernorm
|
||||
--group-query-attention
|
||||
--num-attention-heads 32
|
||||
--num-query-groups 4
|
||||
--kv-channels 128
|
||||
--num-layers 48
|
||||
--hidden-size 2048
|
||||
--ffn-hidden-size 6144
|
||||
|
||||
--normalization RMSNorm
|
||||
--position-embedding-type rope
|
||||
--norm-epsilon 1e-6
|
||||
--rotary-percent 1.0
|
||||
--swiglu
|
||||
--untie-embeddings-and-output-weights
|
||||
--vocab-size 151936
|
||||
|
||||
--rotary-base 1000000
|
||||
|
||||
# moe
|
||||
--moe-ffn-hidden-size 768
|
||||
--moe-router-score-function softmax
|
||||
--moe-token-dispatcher-type alltoall
|
||||
--moe-router-topk 8
|
||||
--moe-layer-freq $MOE_LAYER_FREQ
|
||||
--num-experts 128
|
||||
--moe-grouped-gemm
|
||||
--moe-token-drop-policy probs
|
||||
--moe-router-dtype fp32
|
||||
--moe-permute-fusion
|
||||
--moe-aux-loss-coeff 0
|
||||
)
|
||||
20
configs/models/qwen3-4B.sh
Normal file
20
configs/models/qwen3-4B.sh
Normal file
@@ -0,0 +1,20 @@
|
||||
# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
|
||||
# SPDX-License-Identifier: Apache-2.0
|
||||
|
||||
MODEL_ARGS=(
|
||||
--swiglu
|
||||
--num-layers 36
|
||||
--hidden-size 2560
|
||||
--ffn-hidden-size 9728
|
||||
--num-attention-heads 32
|
||||
--group-query-attention
|
||||
--num-query-groups 8
|
||||
--use-rotary-position-embeddings
|
||||
--disable-bias-linear
|
||||
--normalization "RMSNorm"
|
||||
--norm-epsilon 1e-6
|
||||
--rotary-base "${MODEL_ARGS_ROTARY_BASE:-1000000}"
|
||||
--vocab-size 151936
|
||||
--kv-channels 128
|
||||
--qk-layernorm
|
||||
)
|
||||
21
configs/models/qwen3-8B.sh
Normal file
21
configs/models/qwen3-8B.sh
Normal file
@@ -0,0 +1,21 @@
|
||||
# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
|
||||
# SPDX-License-Identifier: Apache-2.0
|
||||
|
||||
MODEL_ARGS=(
|
||||
--swiglu
|
||||
--num-layers 36
|
||||
--hidden-size 4096
|
||||
--ffn-hidden-size 12288
|
||||
--num-attention-heads 32
|
||||
--group-query-attention
|
||||
--num-query-groups 8
|
||||
--use-rotary-position-embeddings
|
||||
--disable-bias-linear
|
||||
--normalization "RMSNorm"
|
||||
--norm-epsilon 1e-6
|
||||
--rotary-base 1000000
|
||||
--vocab-size 151936
|
||||
--kv-channels 128
|
||||
--qk-layernorm
|
||||
--untie-embeddings-and-output-weights
|
||||
)
|
||||
190
configs/opd/qwen3-4b-opd.py
Normal file
190
configs/opd/qwen3-4b-opd.py
Normal file
@@ -0,0 +1,190 @@
|
||||
# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
|
||||
# SPDX-License-Identifier: Apache-2.0
|
||||
|
||||
import os
|
||||
from pathlib import Path
|
||||
|
||||
import slime.utils.external_utils.command_utils as U
|
||||
from slime.utils.external_utils.command_utils import get_bool_env_var
|
||||
|
||||
# Standard OPD: requires a live teacher server during training.
|
||||
# 2 GPUs for actor (TP=2), 4 GPUs for rollout, 2 GPUs for teacher server.
|
||||
#
|
||||
# Required env vars:
|
||||
# SFT_CHECKPOINT - path to the SFT checkpoint (HF format)
|
||||
|
||||
MODEL_NAME = "Qwen3-4B-Base-Open-Thoughts-Qwen3-8B-sft-3k"
|
||||
TEACHER_MODEL_NAME = "Qwen3-8B"
|
||||
TEACHER_IP = os.environ.get("MASTER_ADDR", "127.0.0.1")
|
||||
TEACHER_PORT = 13141
|
||||
MODEL_TYPE = "qwen3-4B"
|
||||
NUM_GPUS = 8
|
||||
SFT_CHECKPOINT = os.environ["SFT_CHECKPOINT"]
|
||||
|
||||
|
||||
def deploy_teacher_model():
|
||||
import random, string
|
||||
random_suffix = ''.join(random.choices(string.ascii_letters + string.digits, k=6))
|
||||
LOG_FILE = f"/tmp/sglang_{random_suffix}.log"
|
||||
external_ray = get_bool_env_var("SLIME_SCRIPT_EXTERNAL_RAY")
|
||||
|
||||
U.exec_command(
|
||||
"pkill -9 sglang; "
|
||||
"sleep 3; "
|
||||
f"{'' if external_ray else 'ray stop --force; '}"
|
||||
f"{'' if external_ray else 'pkill -9 ray; '}"
|
||||
"pkill -9 slime; "
|
||||
"sleep 3; "
|
||||
f"{'' if external_ray else 'pkill -9 ray; '}"
|
||||
"pkill -9 slime; "
|
||||
"pkill -9 redis; "
|
||||
"true;"
|
||||
)
|
||||
|
||||
U.exec_command(
|
||||
f"CUDA_VISIBLE_DEVICES=6,7 python3 -m sglang.launch_server "
|
||||
f"--model-path /root/models/{TEACHER_MODEL_NAME} "
|
||||
f"--host 0.0.0.0 "
|
||||
f"--port {TEACHER_PORT} "
|
||||
f"--tp 2 "
|
||||
f"--chunked-prefill-size 4096 "
|
||||
f"--mem-fraction-static 0.6 "
|
||||
f"--context-length 32768 "
|
||||
f"> {LOG_FILE} 2>&1 & "
|
||||
)
|
||||
|
||||
U.exec_command(
|
||||
f"until curl -sf http://{TEACHER_IP}:{TEACHER_PORT}/health_generate > /dev/null; do "
|
||||
f" echo 'Waiting for teacher model...'; "
|
||||
f" tail -n 10 {LOG_FILE}; sleep 5; done; "
|
||||
f"echo 'Teacher model ready at {TEACHER_IP}:{TEACHER_PORT}.'; sleep 10;"
|
||||
)
|
||||
|
||||
|
||||
def prepare():
|
||||
U.exec_command("mkdir -p /root/models /root/datasets")
|
||||
U.exec_command(f"huggingface-cli download Qwen/{TEACHER_MODEL_NAME} --local-dir /root/models/{TEACHER_MODEL_NAME}")
|
||||
|
||||
U.convert_checkpoint(
|
||||
model_name=MODEL_NAME,
|
||||
megatron_model_type=MODEL_TYPE,
|
||||
num_gpus_per_node=NUM_GPUS,
|
||||
hf_checkpoint=SFT_CHECKPOINT,
|
||||
)
|
||||
|
||||
deploy_teacher_model()
|
||||
|
||||
|
||||
def execute(rerun=True):
|
||||
load_save_path = f"/root/models/{MODEL_NAME}_ckpt__{Path(__file__).stem}/"
|
||||
|
||||
ckpt_args = (
|
||||
f"--hf-checkpoint {SFT_CHECKPOINT} "
|
||||
f"--ref-load /root/models/{MODEL_NAME}_torch_dist "
|
||||
f"--load {load_save_path} "
|
||||
f"--save {load_save_path} "
|
||||
"--save-interval 10 "
|
||||
"--save-retain-interval 10 "
|
||||
)
|
||||
|
||||
rollout_args = (
|
||||
"--prompt-data /root/datasets/dapo-math-17k/dapo-math-17k.jsonl "
|
||||
"--input-key prompt "
|
||||
"--label-key label "
|
||||
"--apply-chat-template "
|
||||
"--rollout-shuffle "
|
||||
"--num-rollout 3000 "
|
||||
"--rollout-batch-size 64 "
|
||||
"--n-samples-per-prompt 4 "
|
||||
"--rollout-max-response-len 4096 "
|
||||
"--rollout-temperature 0.8 "
|
||||
"--global-batch-size 256 "
|
||||
"--balance-data "
|
||||
)
|
||||
|
||||
rm_args = (
|
||||
"--custom-rm-path slime.rollout.on_policy_distillation.reward_func "
|
||||
"--custom-reward-post-process-path slime.rollout.on_policy_distillation.post_process_rewards "
|
||||
f"--rm-url http://{TEACHER_IP}:{TEACHER_PORT}/generate "
|
||||
"--include-verifiable-reward "
|
||||
)
|
||||
|
||||
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 "
|
||||
"--use-dynamic-batch-size "
|
||||
"--max-tokens-per-gpu 16384 "
|
||||
)
|
||||
|
||||
grpo_args = (
|
||||
"--advantage-estimator on_policy_distillation "
|
||||
"--use-kl-loss "
|
||||
"--kl-loss-coef 0.00 "
|
||||
"--kl-loss-type low_var_kl "
|
||||
"--entropy-coef 0.00 "
|
||||
)
|
||||
|
||||
optimizer_args = (
|
||||
"--optimizer adam "
|
||||
"--lr 2e-6 "
|
||||
"--lr-decay-style constant "
|
||||
"--weight-decay 0.1 "
|
||||
"--adam-beta1 0.9 "
|
||||
"--adam-beta2 0.98 "
|
||||
)
|
||||
|
||||
wandb_args = ""
|
||||
if os.environ.get("WANDB_KEY"):
|
||||
wandb_args = (
|
||||
"--use-wandb "
|
||||
"--wandb-project lightning-opd "
|
||||
f"--wandb-group {Path(__file__).stem} "
|
||||
f"--wandb-key {os.environ['WANDB_KEY']} "
|
||||
)
|
||||
|
||||
sglang_args = (
|
||||
"--rollout-num-gpus-per-engine 1 "
|
||||
"--sglang-mem-fraction-static 0.4 "
|
||||
)
|
||||
|
||||
misc_args = (
|
||||
"--attention-dropout 0.0 "
|
||||
"--hidden-dropout 0.0 "
|
||||
"--accumulate-allreduce-grads-in-fp32 "
|
||||
"--attention-softmax-in-fp32 "
|
||||
"--attention-backend flash "
|
||||
"--actor-num-nodes 1 "
|
||||
"--actor-num-gpus-per-node 2 "
|
||||
"--rollout-num-gpus 4 "
|
||||
)
|
||||
|
||||
train_args = (
|
||||
f"{ckpt_args} "
|
||||
f"{rollout_args} "
|
||||
f"{rm_args} "
|
||||
f"{grpo_args} "
|
||||
f"{optimizer_args} "
|
||||
f"{wandb_args} "
|
||||
f"{perf_args} "
|
||||
f"{sglang_args} "
|
||||
f"{misc_args} "
|
||||
)
|
||||
|
||||
U.execute_train(
|
||||
rerun=rerun,
|
||||
train_args=train_args,
|
||||
num_gpus_per_node=NUM_GPUS,
|
||||
megatron_model_type=MODEL_TYPE,
|
||||
)
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
prepare()
|
||||
execute(rerun=False)
|
||||
190
configs/opd/qwen3-8b-opd.py
Normal file
190
configs/opd/qwen3-8b-opd.py
Normal file
@@ -0,0 +1,190 @@
|
||||
# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
|
||||
# SPDX-License-Identifier: Apache-2.0
|
||||
|
||||
import os
|
||||
from pathlib import Path
|
||||
|
||||
import slime.utils.external_utils.command_utils as U
|
||||
from slime.utils.external_utils.command_utils import get_bool_env_var
|
||||
|
||||
# Standard OPD: requires a live teacher server during training.
|
||||
# 4 GPUs for actor (TP=4), 2 GPUs for rollout, 2 GPUs for teacher server.
|
||||
#
|
||||
# Required env vars:
|
||||
# SFT_CHECKPOINT - path to the SFT checkpoint (HF format)
|
||||
|
||||
MODEL_NAME = "Qwen3-8B-Base-Open-Thoughts-Qwen3-32B-sft-3.5k"
|
||||
TEACHER_MODEL_NAME = "Qwen3-32B"
|
||||
TEACHER_IP = os.environ.get("MASTER_ADDR", "127.0.0.1")
|
||||
TEACHER_PORT = 13141
|
||||
MODEL_TYPE = "qwen3-8B"
|
||||
NUM_GPUS = 8
|
||||
SFT_CHECKPOINT = os.environ["SFT_CHECKPOINT"]
|
||||
|
||||
|
||||
def deploy_teacher_model():
|
||||
import random, string
|
||||
random_suffix = ''.join(random.choices(string.ascii_letters + string.digits, k=6))
|
||||
LOG_FILE = f"/tmp/sglang_{random_suffix}.log"
|
||||
external_ray = get_bool_env_var("SLIME_SCRIPT_EXTERNAL_RAY")
|
||||
|
||||
U.exec_command(
|
||||
"pkill -9 sglang; "
|
||||
"sleep 3; "
|
||||
f"{'' if external_ray else 'ray stop --force; '}"
|
||||
f"{'' if external_ray else 'pkill -9 ray; '}"
|
||||
"pkill -9 slime; "
|
||||
"sleep 3; "
|
||||
f"{'' if external_ray else 'pkill -9 ray; '}"
|
||||
"pkill -9 slime; "
|
||||
"pkill -9 redis; "
|
||||
"true;"
|
||||
)
|
||||
|
||||
U.exec_command(
|
||||
f"CUDA_VISIBLE_DEVICES=6,7 python3 -m sglang.launch_server "
|
||||
f"--model-path /root/models/{TEACHER_MODEL_NAME} "
|
||||
f"--host 0.0.0.0 "
|
||||
f"--port {TEACHER_PORT} "
|
||||
f"--tp 2 "
|
||||
f"--chunked-prefill-size 4096 "
|
||||
f"--mem-fraction-static 0.6 "
|
||||
f"--context-length 32768 "
|
||||
f"> {LOG_FILE} 2>&1 & "
|
||||
)
|
||||
|
||||
U.exec_command(
|
||||
f"until curl -sf http://{TEACHER_IP}:{TEACHER_PORT}/health_generate > /dev/null; do "
|
||||
f" echo 'Waiting for teacher model...'; "
|
||||
f" tail -n 10 {LOG_FILE}; sleep 5; done; "
|
||||
f"echo 'Teacher model ready at {TEACHER_IP}:{TEACHER_PORT}.'; sleep 10;"
|
||||
)
|
||||
|
||||
|
||||
def prepare():
|
||||
U.exec_command("mkdir -p /root/models /root/datasets/dapo-math-17k")
|
||||
U.exec_command(f"huggingface-cli download Qwen/{TEACHER_MODEL_NAME} --local-dir /root/models/{TEACHER_MODEL_NAME}")
|
||||
|
||||
U.convert_checkpoint(
|
||||
model_name=MODEL_NAME,
|
||||
megatron_model_type=MODEL_TYPE,
|
||||
num_gpus_per_node=NUM_GPUS,
|
||||
hf_checkpoint=SFT_CHECKPOINT,
|
||||
)
|
||||
|
||||
deploy_teacher_model()
|
||||
|
||||
|
||||
def execute(rerun=True):
|
||||
load_save_path = f"/root/models/{MODEL_NAME}_ckpt__{Path(__file__).stem}/"
|
||||
|
||||
ckpt_args = (
|
||||
f"--hf-checkpoint {SFT_CHECKPOINT} "
|
||||
f"--ref-load /root/models/{MODEL_NAME}_torch_dist "
|
||||
f"--load {load_save_path} "
|
||||
f"--save {load_save_path} "
|
||||
"--save-interval 5 "
|
||||
"--save-retain-interval 5 "
|
||||
)
|
||||
|
||||
rollout_args = (
|
||||
"--prompt-data /root/datasets/dapo-math-17k/dapo-math-17k.jsonl "
|
||||
"--input-key prompt "
|
||||
"--label-key label "
|
||||
"--apply-chat-template "
|
||||
"--rollout-shuffle "
|
||||
"--num-rollout 3000 "
|
||||
"--rollout-batch-size 64 "
|
||||
"--n-samples-per-prompt 4 "
|
||||
"--rollout-max-response-len 4096 "
|
||||
"--rollout-temperature 0.8 "
|
||||
"--global-batch-size 256 "
|
||||
"--balance-data "
|
||||
)
|
||||
|
||||
rm_args = (
|
||||
"--custom-rm-path slime.rollout.on_policy_distillation.reward_func "
|
||||
"--custom-reward-post-process-path slime.rollout.on_policy_distillation.post_process_rewards "
|
||||
f"--rm-url http://{TEACHER_IP}:{TEACHER_PORT}/generate "
|
||||
"--include-verifiable-reward "
|
||||
)
|
||||
|
||||
perf_args = (
|
||||
"--tensor-model-parallel-size 4 "
|
||||
"--sequence-parallel "
|
||||
"--pipeline-model-parallel-size 1 "
|
||||
"--context-parallel-size 1 "
|
||||
"--expert-model-parallel-size 1 "
|
||||
"--expert-tensor-parallel-size 1 "
|
||||
"--recompute-granularity full "
|
||||
"--recompute-method uniform "
|
||||
"--recompute-num-layers 1 "
|
||||
"--use-dynamic-batch-size "
|
||||
"--max-tokens-per-gpu 8192 "
|
||||
)
|
||||
|
||||
grpo_args = (
|
||||
"--advantage-estimator on_policy_distillation "
|
||||
"--use-kl-loss "
|
||||
"--kl-loss-coef 0.00 "
|
||||
"--kl-loss-type low_var_kl "
|
||||
"--entropy-coef 0.00 "
|
||||
)
|
||||
|
||||
optimizer_args = (
|
||||
"--optimizer adam "
|
||||
"--lr 2e-6 "
|
||||
"--lr-decay-style constant "
|
||||
"--weight-decay 0.1 "
|
||||
"--adam-beta1 0.9 "
|
||||
"--adam-beta2 0.98 "
|
||||
)
|
||||
|
||||
wandb_args = ""
|
||||
if os.environ.get("WANDB_KEY"):
|
||||
wandb_args = (
|
||||
"--use-wandb "
|
||||
"--wandb-project lightning-opd "
|
||||
f"--wandb-group {Path(__file__).stem} "
|
||||
f"--wandb-key {os.environ['WANDB_KEY']} "
|
||||
)
|
||||
|
||||
sglang_args = (
|
||||
"--rollout-num-gpus-per-engine 1 "
|
||||
"--sglang-mem-fraction-static 0.4 "
|
||||
)
|
||||
|
||||
misc_args = (
|
||||
"--attention-dropout 0.0 "
|
||||
"--hidden-dropout 0.0 "
|
||||
"--accumulate-allreduce-grads-in-fp32 "
|
||||
"--attention-softmax-in-fp32 "
|
||||
"--attention-backend flash "
|
||||
"--actor-num-nodes 1 "
|
||||
"--actor-num-gpus-per-node 4 "
|
||||
"--rollout-num-gpus 2 "
|
||||
)
|
||||
|
||||
train_args = (
|
||||
f"{ckpt_args} "
|
||||
f"{rollout_args} "
|
||||
f"{rm_args} "
|
||||
f"{grpo_args} "
|
||||
f"{optimizer_args} "
|
||||
f"{wandb_args} "
|
||||
f"{perf_args} "
|
||||
f"{sglang_args} "
|
||||
f"{misc_args} "
|
||||
)
|
||||
|
||||
U.execute_train(
|
||||
rerun=rerun,
|
||||
train_args=train_args,
|
||||
num_gpus_per_node=NUM_GPUS,
|
||||
megatron_model_type=MODEL_TYPE,
|
||||
)
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
prepare()
|
||||
execute(rerun=False)
|
||||
@@ -0,0 +1,3 @@
|
||||
version https://git-lfs.github.com/spec/v1
|
||||
oid sha256:66cf6cffd11dd523d28cd449afcd5500c62ab47e6617a977039e931d22ae2689
|
||||
size 13944494336
|
||||
23
configs/sft/dataset_info.json
Normal file
23
configs/sft/dataset_info.json
Normal file
@@ -0,0 +1,23 @@
|
||||
{
|
||||
"openthoughts3_300k_qwen3-8b": {
|
||||
"file_name": "data/sft_data/openthoughts3_300k_qwen3-8b.jsonl",
|
||||
"formatting": "sharegpt",
|
||||
"columns": {
|
||||
"messages": "messages"
|
||||
},
|
||||
"tags": {
|
||||
"role_tag": "role",
|
||||
"content_tag": "content",
|
||||
"user_tag": "user",
|
||||
"assistant_tag": "assistant",
|
||||
"system_tag": "system"
|
||||
}
|
||||
},
|
||||
"openthoughts3_300k_qwen3-32b": {
|
||||
"file_name": "data/sft_data/openthoughts3_300k_qwen3-32b.parquet",
|
||||
"formatting": "sharegpt",
|
||||
"columns": {
|
||||
"messages": "messages"
|
||||
}
|
||||
}
|
||||
}
|
||||
16
configs/sft/dataset_info_ori_Fri_5_Jun.json
Normal file
16
configs/sft/dataset_info_ori_Fri_5_Jun.json
Normal file
@@ -0,0 +1,16 @@
|
||||
{
|
||||
"openthoughts3_300k_qwen3-8b": {
|
||||
"file_name": "data/sft_data/openthoughts3_300k_qwen3-8b.jsonl",
|
||||
"formatting": "sharegpt",
|
||||
"columns": {
|
||||
"messages": "messages"
|
||||
}
|
||||
},
|
||||
"openthoughts3_300k_qwen3-32b": {
|
||||
"file_name": "data/sft_data/openthoughts3_300k_qwen3-32b.parquet",
|
||||
"formatting": "sharegpt",
|
||||
"columns": {
|
||||
"messages": "messages"
|
||||
}
|
||||
}
|
||||
}
|
||||
42
configs/sft/qwen3-4b-base-open-thoughts3-qwen3-8b.yaml
Normal file
42
configs/sft/qwen3-4b-base-open-thoughts3-qwen3-8b.yaml
Normal file
@@ -0,0 +1,42 @@
|
||||
### model
|
||||
model_name_or_path: model_weights/qwen3-4b-base
|
||||
|
||||
### method
|
||||
stage: sft
|
||||
do_train: true
|
||||
finetuning_type: full
|
||||
deepspeed: examples/deepspeed/ds_z0_config.json
|
||||
enable_liger_kernel: true
|
||||
packing: true
|
||||
|
||||
|
||||
### dataset
|
||||
dataset: openthoughts3_300k_qwen3-8b
|
||||
template: qwen3
|
||||
cutoff_len: 16384
|
||||
overwrite_cache: true
|
||||
preprocessing_num_workers: 16
|
||||
dataloader_persistent_workers: true
|
||||
dataloader_pin_memory: true
|
||||
dataloader_num_workers: 4
|
||||
|
||||
### output
|
||||
logging_steps: 1
|
||||
save_steps: 100
|
||||
save_total_limit: 10
|
||||
plot_loss: true
|
||||
overwrite_output_dir: false
|
||||
save_only_model: false
|
||||
#report_to: wandb
|
||||
report_to: none
|
||||
run_name: qwen3-4b-base-open-thoughts3-qwen3-8b
|
||||
|
||||
### train
|
||||
per_device_train_batch_size: 4
|
||||
gradient_accumulation_steps: 2
|
||||
learning_rate: 0.00008
|
||||
max_steps: 3000
|
||||
lr_scheduler_type: cosine
|
||||
warmup_ratio: 0.1
|
||||
bf16: true
|
||||
ddp_timeout: 180000000
|
||||
41
configs/sft/qwen3-8b-base-open-thoughts3-qwen3-32b.yaml
Normal file
41
configs/sft/qwen3-8b-base-open-thoughts3-qwen3-32b.yaml
Normal file
@@ -0,0 +1,41 @@
|
||||
### model
|
||||
model_name_or_path: Qwen/Qwen3-8B-Base
|
||||
|
||||
### method
|
||||
stage: sft
|
||||
do_train: true
|
||||
finetuning_type: full
|
||||
deepspeed: examples/deepspeed/ds_z1_config.json
|
||||
enable_liger_kernel: true
|
||||
packing: true
|
||||
|
||||
|
||||
### dataset
|
||||
dataset: openthoughts3_300k_qwen3-32b
|
||||
template: qwen3
|
||||
cutoff_len: 16384
|
||||
overwrite_cache: true
|
||||
preprocessing_num_workers: 16
|
||||
dataloader_persistent_workers: true
|
||||
dataloader_pin_memory: true
|
||||
dataloader_num_workers: 4
|
||||
|
||||
### output
|
||||
logging_steps: 1
|
||||
save_steps: 100
|
||||
save_total_limit: 10
|
||||
plot_loss: true
|
||||
overwrite_output_dir: false
|
||||
save_only_model: false
|
||||
report_to: wandb
|
||||
run_name: qwen3-8b-base-open-thoughts3-qwen3-32b
|
||||
|
||||
### train
|
||||
per_device_train_batch_size: 2
|
||||
gradient_accumulation_steps: 2
|
||||
learning_rate: 8e-5
|
||||
max_steps: 3000
|
||||
lr_scheduler_type: cosine
|
||||
warmup_ratio: 0.1
|
||||
bf16: true
|
||||
ddp_timeout: 180000000
|
||||
49
configs/sft/run_sft.sh
Normal file
49
configs/sft/run_sft.sh
Normal file
@@ -0,0 +1,49 @@
|
||||
#!/usr/bin/env bash
|
||||
# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
|
||||
# SPDX-License-Identifier: Apache-2.0
|
||||
|
||||
# Step 2: Run SFT training with LlamaFactory.
|
||||
#
|
||||
# Required environment variables:
|
||||
# CONFIG_YAML - Name of the SFT config file in configs/sft/ (e.g. qwen3-4b-base-sft-qwen3-8b.yaml)
|
||||
# OUTPUT_DIR - Directory for the SFT checkpoint output
|
||||
#
|
||||
# Optional:
|
||||
# NUM_NODES - Number of nodes (default: 4)
|
||||
# NUM_GPUS - GPUs per node (default: 8)
|
||||
# MASTER_ADDR - Master node address (default: localhost)
|
||||
#
|
||||
# Prerequisites:
|
||||
# - LlamaFactory installed (pip install llamafactory)
|
||||
# - SFT data generated by Step 1 and registered in LlamaFactory's dataset_info.json
|
||||
|
||||
set -euo pipefail
|
||||
|
||||
: "${CONFIG_YAML:?Set CONFIG_YAML (e.g. qwen3-4b-base-sft-qwen3-8b.yaml)}"
|
||||
: "${OUTPUT_DIR:?Set OUTPUT_DIR for SFT checkpoint output}"
|
||||
|
||||
NUM_NODES="${NUM_NODES:-4}"
|
||||
NUM_GPUS="${NUM_GPUS:-8}"
|
||||
MASTER_ADDR="${MASTER_ADDR:-localhost}"
|
||||
MASTER_PORT="${MASTER_PORT:-29500}"
|
||||
|
||||
# torchrun \
|
||||
# --nnodes "${NUM_NODES}" \
|
||||
# --nproc_per_node="${NUM_GPUS}" \
|
||||
# --rdzv_id $RANDOM \
|
||||
# --rdzv_backend c10d \
|
||||
# --rdzv_endpoint "${MASTER_ADDR}:29500" \
|
||||
# -m llamafactory.cli.train \
|
||||
# "configs/sft/${CONFIG_YAML}" \
|
||||
# "dataset_dir=configs/sft" \
|
||||
# "output_dir=${OUTPUT_DIR}"
|
||||
|
||||
FORCE_TORCHRUN=1 \
|
||||
NNODES="${NUM_NODES}" \
|
||||
NPROC_PER_NODE="${NUM_GPUS}" \
|
||||
MASTER_ADDR="${MASTER_ADDR}" \
|
||||
MASTER_PORT="${MASTER_PORT}" \
|
||||
llamafactory-cli train \
|
||||
"configs/sft/${CONFIG_YAML}" \
|
||||
"dataset_dir=configs/sft" \
|
||||
"output_dir=${OUTPUT_DIR}"
|
||||
Reference in New Issue
Block a user