-
Notifications
You must be signed in to change notification settings - Fork 68
[NPU] Add qwen3-4b fully_async script #368
New issue
Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.
By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.
Already on GitHub? Sign in to your account
Merged
+140
−0
Merged
Changes from all commits
Commits
File filter
Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
There are no files selected for viewing
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
| Original file line number | Diff line number | Diff line change | ||||||||||||||||||||||||||||||||
|---|---|---|---|---|---|---|---|---|---|---|---|---|---|---|---|---|---|---|---|---|---|---|---|---|---|---|---|---|---|---|---|---|---|---|
| @@ -0,0 +1,140 @@ | ||||||||||||||||||||||||||||||||||
| #!/bin/bash | ||||||||||||||||||||||||||||||||||
| # Tiny end-to-end fully-async GRPO example using Qwen3-4B-Instruct on the | ||||||||||||||||||||||||||||||||||
| # dapo-math-17k dataset. Designed to run on a single 4-GPU node in a few minutes | ||||||||||||||||||||||||||||||||||
| # | ||||||||||||||||||||||||||||||||||
| # Prerequisites: | ||||||||||||||||||||||||||||||||||
| # /root/models/Qwen3-4B/ (HF checkpoint) | ||||||||||||||||||||||||||||||||||
| # /root/datasets/dapo-math-17k/dapo-math-17k.jsonl | ||||||||||||||||||||||||||||||||||
|
|
||||||||||||||||||||||||||||||||||
| # clean any leftover ray/vllm | ||||||||||||||||||||||||||||||||||
| pkill -9 -f '[v]llm serve|VLL[M]::' | ||||||||||||||||||||||||||||||||||
| pkill -9 -f VLLM | ||||||||||||||||||||||||||||||||||
| sleep 3 | ||||||||||||||||||||||||||||||||||
| ray stop --force | ||||||||||||||||||||||||||||||||||
| pkill -9 ray | ||||||||||||||||||||||||||||||||||
| pkill -9 python | ||||||||||||||||||||||||||||||||||
| sleep 3 | ||||||||||||||||||||||||||||||||||
| pkill -9 ray | ||||||||||||||||||||||||||||||||||
| pkill -9 python | ||||||||||||||||||||||||||||||||||
| pkill -9 redis | ||||||||||||||||||||||||||||||||||
|
|
||||||||||||||||||||||||||||||||||
| set -ex | ||||||||||||||||||||||||||||||||||
|
|
||||||||||||||||||||||||||||||||||
| export PYTHONUNBUFFERED=1 | ||||||||||||||||||||||||||||||||||
| export ASCEND_RT_VISIBLE_DEVICES=0,1,2,3,4,5,6,7,8,9,10,11,12,13,14,15 | ||||||||||||||||||||||||||||||||||
| export RAY_EXPERIMENTAL_NOSET_ASCEND_RT_VISIBLE_DEVICES=1 | ||||||||||||||||||||||||||||||||||
| export RAY_USE_UVLOOP=0 # upstream: Ray's uvloop integration has caused intermittent async actor issues. | ||||||||||||||||||||||||||||||||||
| export CUDA_DEVICE_MAX_CONNECTIONS=1 | ||||||||||||||||||||||||||||||||||
| export HCCL_HOST_SOCKET_PORT_RANGE=60000-60050 | ||||||||||||||||||||||||||||||||||
| export HCCL_NPU_SOCKET_PORT_RANGE=61000-61050 | ||||||||||||||||||||||||||||||||||
| export HYDRA_FULL_ERROR=1 | ||||||||||||||||||||||||||||||||||
| export DISABLE_L2_CACHE=1 | ||||||||||||||||||||||||||||||||||
| export VLLM_ASCEND_ENABLE_NZ=0 | ||||||||||||||||||||||||||||||||||
| export VLLM_USE_AOT_COMPILE=0 | ||||||||||||||||||||||||||||||||||
| export PYTHONPATH="/root/Megatron-Bridge/src:/root/Megatron-LM/:${PYTHONPATH:-}" | ||||||||||||||||||||||||||||||||||
|
|
||||||||||||||||||||||||||||||||||
| unset http_proxy https_proxy HTTP_PROXY HTTPS_PROXY | ||||||||||||||||||||||||||||||||||
|
|
||||||||||||||||||||||||||||||||||
| SCRIPT_DIR="$(cd -- "$(dirname -- "${BASH_SOURCE[0]}")" &>/dev/null && pwd)" | ||||||||||||||||||||||||||||||||||
| source "${SCRIPT_DIR}/../../scripts/models/qwen3-4B.sh" | ||||||||||||||||||||||||||||||||||
|
|
||||||||||||||||||||||||||||||||||
| MODEL_DIR=${MODEL_DIR:-/root/models/Qwen3-4B} | ||||||||||||||||||||||||||||||||||
| DATA_PATH=${DATA_PATH:-/root/datasets/dapo-math-17k/dapo-math-17k.jsonl} | ||||||||||||||||||||||||||||||||||
|
|
||||||||||||||||||||||||||||||||||
| CKPT_ARGS=( | ||||||||||||||||||||||||||||||||||
| --hf-checkpoint "${MODEL_DIR}" | ||||||||||||||||||||||||||||||||||
| --load "${MODEL_DIR}" | ||||||||||||||||||||||||||||||||||
| --ref-load "${MODEL_DIR}" | ||||||||||||||||||||||||||||||||||
| --megatron-to-hf-mode bridge | ||||||||||||||||||||||||||||||||||
| --save /tmp/vime_fully_async_demo/ | ||||||||||||||||||||||||||||||||||
| --save-interval 9999 | ||||||||||||||||||||||||||||||||||
| ) | ||||||||||||||||||||||||||||||||||
|
|
||||||||||||||||||||||||||||||||||
| ROLLOUT_ARGS=( | ||||||||||||||||||||||||||||||||||
| # ↓↓↓ This is the only knob you need to flip to go fully-async ↓↓↓ | ||||||||||||||||||||||||||||||||||
| --rollout-function-path vime.rollout.fully_async_rollout.generate_rollout_fully_async | ||||||||||||||||||||||||||||||||||
|
|
||||||||||||||||||||||||||||||||||
| --prompt-data "${DATA_PATH}" | ||||||||||||||||||||||||||||||||||
| --input-key prompt | ||||||||||||||||||||||||||||||||||
| --label-key label | ||||||||||||||||||||||||||||||||||
| --apply-chat-template | ||||||||||||||||||||||||||||||||||
| --rollout-shuffle | ||||||||||||||||||||||||||||||||||
| --rm-type math | ||||||||||||||||||||||||||||||||||
| --num-rollout 200 | ||||||||||||||||||||||||||||||||||
| --rollout-batch-size 32 | ||||||||||||||||||||||||||||||||||
| --n-samples-per-prompt 8 | ||||||||||||||||||||||||||||||||||
| --rollout-max-response-len 2048 | ||||||||||||||||||||||||||||||||||
| --rollout-temperature 1 | ||||||||||||||||||||||||||||||||||
| --global-batch-size 256 | ||||||||||||||||||||||||||||||||||
| --balance-data | ||||||||||||||||||||||||||||||||||
| ) | ||||||||||||||||||||||||||||||||||
|
|
||||||||||||||||||||||||||||||||||
| 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 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 | ||||||||||||||||||||||||||||||||||
| --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 | ||||||||||||||||||||||||||||||||||
| ) | ||||||||||||||||||||||||||||||||||
|
|
||||||||||||||||||||||||||||||||||
| VLLM_ARGS=( | ||||||||||||||||||||||||||||||||||
| --rollout-num-gpus-per-engine 2 | ||||||||||||||||||||||||||||||||||
| --vllm-gpu-memory-utilization 0.6 | ||||||||||||||||||||||||||||||||||
| ) | ||||||||||||||||||||||||||||||||||
|
|
||||||||||||||||||||||||||||||||||
| MISC_ARGS=( | ||||||||||||||||||||||||||||||||||
| --attention-dropout 0.0 | ||||||||||||||||||||||||||||||||||
| --hidden-dropout 0.0 | ||||||||||||||||||||||||||||||||||
| --accumulate-allreduce-grads-in-fp32 | ||||||||||||||||||||||||||||||||||
| --attention-softmax-in-fp32 | ||||||||||||||||||||||||||||||||||
| --attention-backend flash | ||||||||||||||||||||||||||||||||||
| --use-flash-attn | ||||||||||||||||||||||||||||||||||
| ) | ||||||||||||||||||||||||||||||||||
|
|
||||||||||||||||||||||||||||||||||
| ray start --head --node-ip-address 127.0.0.1 --disable-usage-stats | ||||||||||||||||||||||||||||||||||
|
|
||||||||||||||||||||||||||||||||||
| # fully-async splits actor / rollout onto disjoint GPUs (no colocation). | ||||||||||||||||||||||||||||||||||
| ACTOR_GPUS=2 | ||||||||||||||||||||||||||||||||||
| ROLLOUT_GPUS=2 | ||||||||||||||||||||||||||||||||||
|
|
||||||||||||||||||||||||||||||||||
| ray job submit --address="http://127.0.0.1:8265" \ | ||||||||||||||||||||||||||||||||||
| -- python3 train_async.py \ | ||||||||||||||||||||||||||||||||||
| --actor-num-nodes 1 \ | ||||||||||||||||||||||||||||||||||
| --actor-num-gpus-per-node "${ACTOR_GPUS}" \ | ||||||||||||||||||||||||||||||||||
| --rollout-num-gpus "${ROLLOUT_GPUS}" \ | ||||||||||||||||||||||||||||||||||
| ${MODEL_ARGS[@]} \ | ||||||||||||||||||||||||||||||||||
| ${CKPT_ARGS[@]} \ | ||||||||||||||||||||||||||||||||||
| ${ROLLOUT_ARGS[@]} \ | ||||||||||||||||||||||||||||||||||
| ${OPTIMIZER_ARGS[@]} \ | ||||||||||||||||||||||||||||||||||
| ${GRPO_ARGS[@]} \ | ||||||||||||||||||||||||||||||||||
| ${PERF_ARGS[@]} \ | ||||||||||||||||||||||||||||||||||
| ${VLLM_ARGS[@]} \ | ||||||||||||||||||||||||||||||||||
| ${MISC_ARGS[@]} | ||||||||||||||||||||||||||||||||||
|
Comment on lines
+133
to
+140
There was a problem hiding this comment. Choose a reason for hiding this commentThe reason will be displayed to describe this comment to others. Learn more. Array expansions in Bash should be double-quoted (e.g.,
Suggested change
|
||||||||||||||||||||||||||||||||||
Add this suggestion to a batch that can be applied as a single commit.
This suggestion is invalid because no changes were made to the code.
Suggestions cannot be applied while the pull request is closed.
Suggestions cannot be applied while viewing a subset of changes.
Only one suggestion per line can be applied in a batch.
Add this suggestion to a batch that can be applied as a single commit.
Applying suggestions on deleted lines is not supported.
You must change the existing code in this line in order to create a valid suggestion.
Outdated suggestions cannot be applied.
This suggestion has been applied or marked resolved.
Suggestions cannot be applied from pending reviews.
Suggestions cannot be applied on multi-line comments.
Suggestions cannot be applied while the pull request is queued to merge.
Suggestion cannot be applied right now. Please check back later.
There was a problem hiding this comment.
Choose a reason for hiding this comment
The reason will be displayed to describe this comment to others. Learn more.
The script header states it is designed to run on a single 4-GPU node, but
ASCEND_RT_VISIBLE_DEVICESis configured with 16 devices (0-15). It should be restricted to0,1,2,3to align with the 4-GPU design and avoid potential runtime issues on 4-GPU systems.