Skip to content

Commit 1bc1205

Browse files
UsernameFullUsernameFull
committed
Squash npu_fp8 onto alibaba/main
Consolidates the committed npu_fp8 changes into one commit.
1 parent 370cb24 commit 1bc1205

43 files changed

Lines changed: 2311 additions & 189 deletions

Some content is hidden

Large Commits have some content hidden by default. Use the searchbox below for content that may be hidden.

‎docker/Dockerfile.A5‎

Lines changed: 75 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,75 @@
1+
FROM swr.cn-south-1.myhuaweicloud.com/ascendhub/cann:9.1.0-950-ubuntu22.04-py3.12-devel
2+
3+
ARG VLLM_VERSION=v0.23.0
4+
ARG VLLM_ASCEND_VERSION=v0.23.0rc1
5+
ARG MEGATRON_VERSION=core_r0.17.0
6+
ARG MEGATRON_ADAPTOR_VERSION=core_r0.17.0
7+
8+
ENV DEBIAN_FRONTEND=noninteractive \
9+
PIP_INDEX_URL=https://repo.huaweicloud.com/repository/pypi/simple \
10+
VLLM_TARGET_DEVICE=empty \
11+
ASCEND_HOME_PATH=/usr/local/Ascend/ascend-toolkit/latest \
12+
SOC_VERSION=ascend950pr_9599 \
13+
HF_ENDPOINT=https://hf-mirror.com \
14+
VLLM_USE_MODELSCOPE=True \
15+
PYTORCH_NPU_ALLOC_CONF=expandable_segments:True \
16+
TASK_QUEUE_ENABLE=2 \
17+
HCCL_NPU_SOCKET_PORT_RANGE=auto
18+
19+
SHELL ["/bin/bash", "-c"]
20+
21+
RUN apt-get update && \
22+
apt-get install -y --no-install-recommends \
23+
build-essential ca-certificates git iproute2 \
24+
openjdk-17-jdk pkg-config tmux zip && \
25+
rm -rf /var/lib/apt/lists/*
26+
27+
RUN python3 -m pip install --upgrade pip setuptools wheel setuptools-rust setuptools-scm packaging "cmake>=3.26" ninja && \
28+
python3 -m pip install --index-url https://download.pytorch.org/whl/cpu \
29+
torch==2.10.0 torchvision==0.25.0 torchaudio==2.10.0 && \
30+
python3 -m pip install --no-deps torch-npu==2.10.0.post4 && \
31+
python3 -m pip install --extra-index-url https://triton-ascend.osinfra.cn/pypi/simple \
32+
triton-ascend==3.2.2
33+
34+
WORKDIR /workspace
35+
36+
RUN git clone --depth 1 --branch ${VLLM_VERSION} \
37+
https://github.com/vllm-project/vllm.git vllm-${VLLM_VERSION} && \
38+
git clone --depth 1 --branch ${VLLM_ASCEND_VERSION} \
39+
https://github.com/vllm-project/vllm-ascend.git vllm-ascend-${VLLM_ASCEND_VERSION} && \
40+
git clone --depth 1 --branch ${MEGATRON_VERSION} \
41+
https://github.com/NVIDIA/Megatron-LM.git Megatron-LM && \
42+
git clone --depth 1 https://gitcode.com/Ascend/TransformerEngineNPU TransformerEngineNPU && \
43+
git clone --depth 1 --branch ${MEGATRON_ADAPTOR_VERSION} \
44+
https://gitcode.com/Ascend/MegatronAdaptor MegatronAdaptor
45+
46+
RUN source /usr/local/Ascend/ascend-toolkit/latest/set_env.sh && \
47+
source /usr/local/Ascend/nnal/atb/set_env.sh && \
48+
python3 -m pip install -r /workspace/vllm-${VLLM_VERSION}/requirements/common.txt && \
49+
python3 -m pip install numpy==1.26.4 opencv-python-headless==4.10.0.84 && \
50+
python3 -m pip install --no-build-isolation --no-deps -e /workspace/vllm-${VLLM_VERSION} && \
51+
python3 -m pip install decorator msgpack numba pandas pandas-stubs \
52+
pybind11 quart scipy && \
53+
python3 -m pip install --no-build-isolation --no-deps \
54+
-e /workspace/vllm-ascend-${VLLM_ASCEND_VERSION}
55+
56+
WORKDIR /workspace/ROLL
57+
COPY . .
58+
59+
RUN source /usr/local/Ascend/ascend-toolkit/latest/set_env.sh && \
60+
source /usr/local/Ascend/nnal/atb/set_env.sh && \
61+
python3 -m pip install --ignore-requires-python gem-llm==0.0.4 && \
62+
grep -v '^gem-llm==0.0.4$' requirements_common.txt > /tmp/requirements_common_nogem.txt && \
63+
python3 -m pip install -r /tmp/requirements_common_nogem.txt && \
64+
python3 -m pip install deepspeed==0.16.4 tensorboard && \
65+
python3 -m pip install --no-build-isolation -e /workspace/Megatron-LM && \
66+
python3 -m pip install --no-build-isolation -e /workspace/TransformerEngineNPU && \
67+
python3 -m pip install --no-build-isolation -e /workspace/MegatronAdaptor && \
68+
python3 -m pip install --no-build-isolation -e /workspace/ROLL && \
69+
rm -f /tmp/requirements_common_nogem.txt
70+
71+
RUN echo "source /usr/local/Ascend/ascend-toolkit/set_env.sh" >> /root/.bashrc && \
72+
echo "source /usr/local/Ascend/nnal/atb/set_env.sh" >> /root/.bashrc && \
73+
echo "export HCCL_NPU_SOCKET_PORT_RANGE=auto" >> /root/.bashrc
74+
75+
CMD ["bash"]

‎examples/ascend_examples/qwen3_0_6b_rlvr_megatron_bf16_bf16_4npu.yaml‎

Lines changed: 152 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,152 @@
1+
hydra:
2+
run:
3+
dir: .
4+
output_subdir: null
5+
6+
exp_name: "qwen3-0.6B-rlvr-megatron-fp8-vllm-mxfp8-4npu"
7+
seed: 42
8+
logging_dir: ./output/logs
9+
output_dir: ./output
10+
system_envs:
11+
USE_MODELSCOPE: '1'
12+
13+
checkpoint_config:
14+
type: file_system
15+
output_dir: ./rl_examples/models/${exp_name}
16+
17+
track_with: tensorboard
18+
tracker_kwargs:
19+
log_dir: ./rl_examples/llm/tensorboard/roll_exp/rlvr
20+
rpc_timeout: 72000
21+
22+
num_gpus_per_node: 4
23+
24+
max_steps: 500
25+
save_steps: 100
26+
logging_steps: 1
27+
eval_steps: 10
28+
resume_from_checkpoint: false
29+
30+
rollout_batch_size: 32
31+
prompt_length: 2048
32+
response_length: 4096
33+
34+
num_return_sequences_in_group: 8
35+
ppo_epochs: 1
36+
adv_estimator: "reinforce"
37+
38+
value_clip: 0.5
39+
reward_clip: 10
40+
advantage_clip: 2.0
41+
dual_clip_loss: true
42+
43+
norm_mean_type: ~
44+
norm_std_type: ~
45+
46+
max_len_mask: true
47+
difficulty_mask: true
48+
difficulty_low_threshold: 0.1
49+
difficulty_high_threshold: 0.95
50+
error_max_len_clip: false
51+
52+
difficulty_loss_weight: false
53+
length_loss_weight: false
54+
55+
add_token_level_kl: false
56+
whiten_advantages: true
57+
58+
pretrain: Qwen/Qwen3-0.6B
59+
reward_pretrain: Qwen/Qwen3-0.6B
60+
61+
actor_train:
62+
model_args:
63+
disable_gradient_checkpointing: false
64+
dtype: bf16
65+
model_type: ~
66+
training_args:
67+
learning_rate: 1.0e-6
68+
weight_decay: 0
69+
per_device_train_batch_size: 2
70+
gradient_accumulation_steps: 8
71+
warmup_steps: 20
72+
num_train_epochs: 50
73+
data_args:
74+
template: qwen2_5
75+
file_name:
76+
- data/math_deepmath_deal.jsonl
77+
domain_interleave_probs:
78+
math_rule: 1
79+
dataset_dir: data
80+
messages: messages
81+
interleave_probs: "1.0"
82+
preprocessing_num_workers: 16
83+
strategy_args:
84+
strategy_name: megatron_train
85+
strategy_config:
86+
tensor_model_parallel_size: 2
87+
pipeline_model_parallel_size: 1
88+
expert_model_parallel_size: 1
89+
expert_tensor_parallel_size: 1
90+
use_distributed_optimizer: true
91+
recompute_granularity: full
92+
recompute_method: uniform
93+
recompute_num_layers: 1
94+
use_flash_attn: true
95+
moe_router_dtype: fp32
96+
use_sequence_packing: false
97+
use_remove_padding: false
98+
device_mapping: list(range(0,4))
99+
infer_batch_size: 2
100+
101+
actor_infer:
102+
model_args:
103+
disable_gradient_checkpointing: true
104+
dtype: bf16
105+
generating_args:
106+
max_new_tokens: ${response_length}
107+
top_p: 0.99
108+
top_k: 100
109+
num_beams: 1
110+
temperature: 0.99
111+
num_return_sequences: ${num_return_sequences_in_group}
112+
data_args:
113+
template: qwen2_5
114+
strategy_args:
115+
strategy_name: vllm
116+
strategy_config:
117+
gpu_memory_utilization: 0.8
118+
block_size: 16
119+
max_model_len: 6144
120+
tensor_parallel_size: 2
121+
enforce_eager: true
122+
load_format: dummy
123+
device_mapping: list(range(0,4))
124+
infer_batch_size: 1
125+
126+
reference:
127+
model_args:
128+
disable_gradient_checkpointing: true
129+
dtype: bf16
130+
model_type: ~
131+
data_args:
132+
template: qwen2_5
133+
strategy_args:
134+
strategy_name: megatron_infer
135+
strategy_config:
136+
tensor_model_parallel_size: 2
137+
pipeline_model_parallel_size: 1
138+
expert_model_parallel_size: 1
139+
expert_tensor_parallel_size: 1
140+
device_mapping: list(range(0,4))
141+
infer_batch_size: 2
142+
143+
rewards:
144+
math_rule:
145+
worker_cls: roll.pipeline.rlvr.rewards.math_rule_reward_worker.MathRuleRewardWorker
146+
model_args:
147+
model_name_or_path: ${reward_pretrain}
148+
data_args:
149+
template: qwen2_5
150+
tag_included: [deepmath_103k, aime]
151+
world_size: 4
152+
infer_batch_size: 1

‎examples/ascend_examples/qwen3_0_6b_rlvr_megatron_fp8_mxfp8_4npu.yaml‎

Lines changed: 157 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,157 @@
1+
hydra:
2+
run:
3+
dir: .
4+
output_subdir: null
5+
6+
exp_name: "qwen3-0.6B-rlvr-megatron-fp8-vllm-mxfp8-4npu"
7+
seed: 42
8+
logging_dir: ./output/logs
9+
output_dir: ./output
10+
system_envs:
11+
USE_MODELSCOPE: '1'
12+
13+
checkpoint_config:
14+
type: file_system
15+
output_dir: ./rl_examples/models/${exp_name}
16+
17+
track_with: tensorboard
18+
tracker_kwargs:
19+
log_dir: ./rl_examples/llm/tensorboard/roll_exp/rlvr
20+
rpc_timeout: 72000
21+
22+
num_gpus_per_node: 4
23+
24+
max_steps: 500
25+
save_steps: 100
26+
logging_steps: 1
27+
eval_steps: 10
28+
resume_from_checkpoint: false
29+
30+
rollout_batch_size: 32
31+
prompt_length: 2048
32+
response_length: 4096
33+
34+
num_return_sequences_in_group: 8
35+
ppo_epochs: 1
36+
adv_estimator: "reinforce"
37+
38+
value_clip: 0.5
39+
reward_clip: 10
40+
advantage_clip: 2.0
41+
dual_clip_loss: true
42+
43+
norm_mean_type: ~
44+
norm_std_type: ~
45+
46+
max_len_mask: true
47+
difficulty_mask: true
48+
difficulty_low_threshold: 0.1
49+
difficulty_high_threshold: 0.95
50+
error_max_len_clip: false
51+
52+
difficulty_loss_weight: false
53+
length_loss_weight: false
54+
55+
add_token_level_kl: false
56+
whiten_advantages: true
57+
58+
pretrain: Qwen/Qwen3-0.6B
59+
reward_pretrain: Qwen/Qwen3-0.6B
60+
61+
actor_train:
62+
model_args:
63+
disable_gradient_checkpointing: false
64+
dtype: bf16
65+
model_type: ~
66+
training_args:
67+
learning_rate: 1.0e-6
68+
weight_decay: 0
69+
per_device_train_batch_size: 2
70+
gradient_accumulation_steps: 8
71+
warmup_steps: 20
72+
num_train_epochs: 50
73+
data_args:
74+
template: qwen2_5
75+
file_name:
76+
- data/math_deepmath_deal.jsonl
77+
domain_interleave_probs:
78+
math_rule: 1
79+
dataset_dir: data
80+
messages: messages
81+
interleave_probs: "1.0"
82+
preprocessing_num_workers: 16
83+
strategy_args:
84+
strategy_name: megatron_train
85+
strategy_config:
86+
tensor_model_parallel_size: 2
87+
pipeline_model_parallel_size: 1
88+
expert_model_parallel_size: 1
89+
expert_tensor_parallel_size: 1
90+
use_distributed_optimizer: true
91+
recompute_granularity: full
92+
recompute_method: uniform
93+
recompute_num_layers: 1
94+
fp8: e4m3
95+
fp8_recipe: mxfp8
96+
use_flash_attn: true
97+
moe_router_dtype: fp32
98+
use_sequence_packing: false
99+
use_remove_padding: false
100+
device_mapping: list(range(0,4))
101+
infer_batch_size: 2
102+
103+
actor_infer:
104+
model_args:
105+
disable_gradient_checkpointing: true
106+
dtype: bf16
107+
generating_args:
108+
max_new_tokens: ${response_length}
109+
top_p: 0.99
110+
top_k: 100
111+
num_beams: 1
112+
temperature: 0.99
113+
num_return_sequences: ${num_return_sequences_in_group}
114+
data_args:
115+
template: qwen2_5
116+
strategy_args:
117+
strategy_name: vllm
118+
strategy_config:
119+
gpu_memory_utilization: 0.8
120+
block_size: 16
121+
max_model_len: 6144
122+
tensor_parallel_size: 2
123+
enforce_eager: true
124+
load_format: dummy
125+
online_quantization: ascend_mxfp8
126+
online_quantization_config:
127+
group_size: 32
128+
device_mapping: list(range(0,4))
129+
infer_batch_size: 1
130+
131+
reference:
132+
model_args:
133+
disable_gradient_checkpointing: true
134+
dtype: bf16
135+
model_type: ~
136+
data_args:
137+
template: qwen2_5
138+
strategy_args:
139+
strategy_name: megatron_infer
140+
strategy_config:
141+
tensor_model_parallel_size: 2
142+
pipeline_model_parallel_size: 1
143+
expert_model_parallel_size: 1
144+
expert_tensor_parallel_size: 1
145+
device_mapping: list(range(0,4))
146+
infer_batch_size: 2
147+
148+
rewards:
149+
math_rule:
150+
worker_cls: roll.pipeline.rlvr.rewards.math_rule_reward_worker.MathRuleRewardWorker
151+
model_args:
152+
model_name_or_path: ${reward_pretrain}
153+
data_args:
154+
template: qwen2_5
155+
tag_included: [deepmath_103k, aime]
156+
world_size: 4
157+
infer_batch_size: 1

0 commit comments

Comments
 (0)