Skip to content

Commit df340d7

Browse files
ArronHZGArronHZG
authored
[megatron] feat: Using MTP in RL Training and Inference (verl-project#4936)
# Guide to Using MTP in RL Training and Inference **Author**: `https://github.com/meituan-search` **Last Updated**: 2026/01/14 # 1. Scope of Support Currently, RL training can be performed on mimo-7B-RL, Qwen-next, and Deepseek series models based on the MTP architecture. The support rules for training and inference engines are as follows: - **Training Engine**: Only supports the `mbridge + megatron` combination; other training engines are not compatible at this time; - **Inference Engine**: Compatible with all engines, but the model must be in the corresponding engine's compatibility list; - **Dependency Versions**: - mbridge: Use the specified branch: [https://github.com/ArronHZG/mbridge/tree/feature/verl_mtp](https://github.com/ArronHZG/mbridge/tree/feature/verl_mtp) (will be merged into the main branch in the future); - megatron: Use the latest dev version (commit: [23e092f41ec8bc659020e401ddac9576c1cfed7e](https://github.com/NVIDIA/Megatron-LM/tree/23e092f41ec8bc659020e401ddac9576c1cfed7e)), which supports MTP + CP training methods. # 2. MTP Training Configuration (Core Parameters) The MTP training process can be flexibly controlled through the following configurations. All configurations are based on the `actor_rollout_ref.model.mtp` prefix: | Configuration Scenario | Core Parameters | Description | |------------------------|---------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------|---------------------------------------------------------| | Load MTP Parameters Only | `enable=True` | VRAM usage will increase, but the exported parameters include the MTP module and can be directly used for online deployment | | Full-Parameter MTP Training | `enable=True`<br>`enable_train=True`<br>`mtp_loss_scaling_factor=0.1` | MTP Loss will apply to all model parameters | | MTP Parameter-Only Training | `enable=True`<br>`enable_train=True`<br>`detach_encoder=True` | Freeze the Encoder layer, update only MTP module parameters, MTP Loss applies only to MTP parameters | | MTP Accelerated Rollout | 1. vLLM configuration:<br>`enable=True`<br>`enable_rollout=True`<br>`method="mtp"`<br>`num_speculative_tokens=1`<br>2. SGLang configuration:<br>`enable=True`<br>`enable_rollout=True`<br>`speculative_algorithm="EAGLE"`<br>`speculative_num_steps=2`<br>`speculative_eagle_topk=2`<br>`speculative_num_draft_tokens=4` | Achieve inference acceleration during the Rollout phase based on MTP | # 3. Experimental Results The experiment was conducted as follows: * model = mimo-7B-math * max_response_length = 8k Experiment chart: ![fully_async_policy_revenue]( https://github.com/ArronHZG/verl-community/blob/main/docs/mimo-7b-mtp.png?raw=true) **Scenarios with No Significant Effect** The following configurations will not have a noticeable impact on training results: 1. The base model does not carry MTP parameters; 2. The base model carries MTP parameters, but the MTP module is not trained; 3. The base model carries MTP parameters and trains MTP, with `mtp_loss_scaling_factor=0`; 4. The base model carries MTP parameters, trains MTP and detaches the encoder, with `mtp_loss_scaling_factor=0.1`. **Scenarios with Significant Effect** Only the following configuration will have a noticeable impact on training results: - The base model carries MTP parameters, MTP Loss applies to all model parameters, and `mtp_loss_scaling_factor=0.1`. **Recommended Training Method** It is recommended to adopt the `detach_encoder=True` approach for MTP training. # 4. Performance Notes for MTP in Rollout Inference The effectiveness of MTP-accelerated Rollout is significantly affected by **model size** and **inference hardware**. Key reference information is as follows: **Hardware Tensor Core Performance** | Hardware Model | FP16 Performance (TFLOPS) | |----------------|---------------------------| | H20 | 148 | | H800 | 1,671 | | H200 | 1,979 | **Measured Performance and Recommendations** Taking the mimo-7B model deployed separately on H20 hardware using SGLang as an example: After enabling MTP speculative decoding, the Rollout throughput decreases by approximately 50%. - Current priority recommendation: Do not enable MTP acceleration during the inference phase for now; - Future planning: Further optimization of the speculative logic in the Rollout phase will be conducted to improve throughput performance. > Add **concise** overview of what this PR aims to achieve or accomplish. Reference related GitHub issues and PRs that help with the review. ### Checklist Before Starting - [x] Search for similar PRs. Paste at least one query link here: ... - [x] Format the PR title as `[{modules}] {type}: {description}` (This will be checked by the CI) - `{modules}` include `fsdp`, `megatron`, `veomni`, `sglang`, `vllm`, `rollout`, `trainer`, `ci`, `training_utils`, `recipe`, `hardware`, `deployment`, `ray`, `worker`, `single_controller`, `misc`, `perf`, `model`, `algo`, `env`, `tool`, `ckpt`, `doc`, `data`, `cfg`, `reward` - If this PR involves multiple modules, separate them with `,` like `[megatron, fsdp, doc]` - `{type}` is in `feat`, `fix`, `refactor`, `chore`, `test` - If this PR breaks any API (CLI arguments, config, function signature, etc.), add `[BREAKING]` to the beginning of the title. - Example: `[BREAKING][fsdp, megatron] feat: dynamic batching` ### Test > For changes that can not be tested by CI (e.g., algorithm implementation, new model support), validate by experiment(s) and show results like training curve plots, evaluation results, etc. ### API and Usage Example > Demonstrate how the API changes if any, and provide usage example(s) if possible. ```python # Add code snippet or script demonstrating how to use this ``` ### Design & Code Changes > Demonstrate the high-level design if this PR is complex, and list the specific changes. ### Checklist Before Submitting > [!IMPORTANT] > Please check all the following items before requesting a review, otherwise the reviewer might deprioritize this PR for review. - [x] Read the [Contribute Guide](https://github.com/volcengine/verl/blob/main/CONTRIBUTING.md). - [x] Apply [pre-commit checks](https://github.com/volcengine/verl/blob/main/CONTRIBUTING.md#code-linting-and-formatting): `pre-commit install && pre-commit run --all-files --show-diff-on-failure --color=always` - [x] Add / Update [the documentation](https://github.com/volcengine/verl/tree/main/docs). - [x] Add unit or end-to-end test(s) to [the CI workflow](https://github.com/volcengine/verl/tree/main/.github/workflows) to cover all the code. If not feasible, explain why: ... - [x] Once your PR is ready for CI, send a message in [the `ci-request` channel](https://verl-project.slack.com/archives/C091TCESWB1) in [the `verl` Slack workspace](https://join.slack.com/t/verl-project/shared_invite/zt-3855yhg8g-CTkqXu~hKojPCmo7k_yXTQ). (If not accessible, please try [the Feishu group (飞书群)](https://applink.larkoffice.com/client/chat/chatter/add_by_link?link_token=772jd4f1-cd91-441e-a820-498c6614126a).) - [x] If your PR is related to the `recipe` submodule, please also update the reference to the submodule commit via `git submodule update --remote` or `cd recipe && git pull origin main`.
1 parent f104dfa commit df340d7

File tree

23 files changed

+778
-15
lines changed

23 files changed

+778
-15
lines changed

‎docs/advance/fully_async.md‎

Lines changed: 3 additions & 3 deletions
Original file line numberDiff line numberDiff line change
@@ -63,7 +63,7 @@ Currently, the supported usage mode is Megatron/FSDP+vLLM/SGLang. vLLM/SGLang mu
6363
The overall architecture of fully_async_policy is shown in the figure below. fully_async_policy mainly consists of four
6464
parts: Rollouter, MessageQueue, Trainer, and ParameterSynchronizer.
6565

66-
![fully_async_policy_structure](https://github.com/ArronHZG/verl-community/blob/recipe/async_policy/docs/fully_async_policy_structure.svg?raw=true)
66+
![fully_async_policy_structure](https://github.com/ArronHZG/verl-community/blob/main/docs/fully_async_policy_structure.svg?raw=true)
6767

6868
1. Rollouter generates sequences sample by sample and puts the generated samples into the MessageQueue, with the
6969
production speed controlled by freshness.
@@ -79,7 +79,7 @@ After we perform resource isolation, the time for rollout and train may be longe
7979
are used),
8080
but the overlap in their time consumption reduces the end-to-end time consumption.
8181

82-
![fully_async_policy_revenue](https://github.com/ArronHZG/verl-community/blob/recipe/async_policy/docs/fully_async_policy_revenue.svg?raw=true)
82+
![fully_async_policy_revenue](https://github.com/ArronHZG/verl-community/blob/main/docs/fully_async_policy_revenue.svg?raw=true)
8383

8484
## Usage
8585

@@ -246,7 +246,7 @@ but the overlap in their time consumption reduces the end-to-end time consumptio
246246
generated after synchronization. This reduces the time to wait for active tasks to finish.
247247
3. As shown in figure d;
248248

249-
![fully_async_policy_mode](https://github.com/ArronHZG/verl-community/blob/recipe/async_policy/docs/fully_async_policy_mode.svg?raw=true)
249+
![fully_async_policy_mode](https://github.com/ArronHZG/verl-community/blob/main/docs/fully_async_policy_mode.svg?raw=true)
250250

251251
### Key Metrics
252252

‎docs/advance/mtp.md‎

Lines changed: 84 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,84 @@
1+
# Guide to Using MTP in RL Training and Inference
2+
3+
**Author**: `https://github.com/meituan-search`
4+
5+
Last updated: 01/16/2026
6+
7+
# 1. Scope of Support
8+
9+
Currently, RL training can be performed on mimo-7B-RL, Qwen-next, and Deepseek series models based on the MTP architecture. The support rules for training and inference engines are as follows:
10+
11+
- **Training Engine**: Only supports the `mbridge + megatron` combination; other training engines are not compatible at this time;
12+
13+
- **Inference Engine**: Compatible with all engines, but the model must be in the corresponding engine's compatibility list;
14+
15+
- **Dependency Versions**:
16+
17+
- mbridge: Use the specified branch: [https://github.com/ArronHZG/mbridge/tree/feature/verl_mtp](https://github.com/ArronHZG/mbridge/tree/feature/verl_mtp) (will be merged into the main branch in the future);
18+
19+
- megatron: Use the latest dev version (commit: [23e092f41ec8bc659020e401ddac9576c1cfed7e](https://github.com/NVIDIA/Megatron-LM/tree/23e092f41ec8bc659020e401ddac9576c1cfed7e)), which supports MTP + CP training methods.
20+
21+
# 2. MTP Training Configuration (Core Parameters)
22+
23+
The MTP training process can be flexibly controlled through the following configurations. All configurations are based on the `actor_rollout_ref.model.mtp` prefix:
24+
25+
| Configuration Scenario | Core Parameters | Description |
26+
|------------------------|---------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------|---------------------------------------------------------|
27+
| Load MTP Parameters Only | `enable=True` | VRAM usage will increase, but the exported parameters include the MTP module and can be directly used for online deployment |
28+
| Full-Parameter MTP Training | `enable=True`<br>`enable_train=True`<br>`mtp_loss_scaling_factor=0.1` | MTP Loss will apply to all model parameters |
29+
| MTP Parameter-Only Training | `enable=True`<br>`enable_train=True`<br>`detach_encoder=True` | Freeze the Encoder layer, update only MTP module parameters, MTP Loss applies only to MTP parameters |
30+
| MTP Accelerated Rollout | 1. vLLM configuration:<br>`enable=True`<br>`enable_rollout=True`<br>`method="mtp"`<br>`num_speculative_tokens=1`<br>2. SGLang configuration:<br>`enable=True`<br>`enable_rollout=True`<br>`speculative_algorithm="EAGLE"`<br>`speculative_num_steps=2`<br>`speculative_eagle_topk=2`<br>`speculative_num_draft_tokens=4` | Achieve inference acceleration during the Rollout phase based on MTP |
31+
32+
# 3. Experimental Results
33+
34+
The experiment was conducted as follows:
35+
36+
* model = mimo-7B-math
37+
* max_response_length = 8k
38+
39+
Experiment chart:
40+
41+
![fully_async_policy_revenue](
42+
https://github.com/ArronHZG/verl-community/blob/main/docs/mimo-7b-mtp.png?raw=true)
43+
44+
**Scenarios with No Significant Effect**
45+
46+
The following configurations will not have a noticeable impact on training results:
47+
48+
1. The base model does not carry MTP parameters;
49+
50+
2. The base model carries MTP parameters, but the MTP module is not trained;
51+
52+
3. The base model carries MTP parameters and trains MTP, with `mtp_loss_scaling_factor=0`;
53+
54+
4. The base model carries MTP parameters, trains MTP and detaches the encoder, with `mtp_loss_scaling_factor=0.1`.
55+
56+
**Scenarios with Significant Effect**
57+
58+
Only the following configuration will have a noticeable impact on training results:
59+
60+
- The base model carries MTP parameters, MTP Loss applies to all model parameters, and `mtp_loss_scaling_factor=0.1`.
61+
62+
**Recommended Training Method**
63+
64+
It is recommended to adopt the `detach_encoder=True` approach for MTP training.
65+
66+
# 4. Performance Notes for MTP in Rollout Inference
67+
68+
The effectiveness of MTP-accelerated Rollout is significantly affected by **model size** and **inference hardware**. Key reference information is as follows:
69+
70+
**Hardware Tensor Core Performance**
71+
72+
| Hardware Model | FP16 Performance (TFLOPS) |
73+
|----------------|---------------------------|
74+
| H20 | 148 |
75+
| H800 | 1,671 |
76+
| H200 | 1,979 |
77+
78+
**Measured Performance and Recommendations**
79+
80+
Taking the mimo-7B model deployed separately on H20 hardware using SGLang as an example: After enabling MTP speculative decoding, the Rollout throughput decreases by approximately 50%.
81+
82+
- Current priority recommendation: Do not enable MTP acceleration during the inference phase for now;
83+
84+
- Future planning: Further optimization of the speculative logic in the Rollout phase will be conducted to improve throughput performance.

‎docs/index.rst‎

Lines changed: 1 addition & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -134,6 +134,7 @@ verl is fast with:
134134
advance/grafana_prometheus.md
135135
advance/fp8.md
136136
advance/async-on-policy-distill
137+
advance/mtp.md
137138

138139
.. toctree::
139140
:maxdepth: 1

‎examples/mtp_trainer/runtime_env.yaml‎

Lines changed: 16 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,16 @@
1+
working_dir: ./
2+
3+
excludes:
4+
- ".git/"
5+
6+
env_vars:
7+
VLLM_USE_V1: "1"
8+
HYDRA_FULL_ERROR: "1"
9+
NCCL_NVLS_ENABLE: "0"
10+
NCCL_SOCKET_IFNAME: "eth0"
11+
TMPDIR: "/tmp"
12+
CUDA_HOME: "/usr/local/cuda"
13+
CUDA_TMPDIR: "/tmp"
14+
CUDA_CACHE_PATH: "/tmp/cuda_cache"
15+
HF_HOME: "/tmp/hf_home_mimo"
16+
PYTHONPATH: "/tmp/hf_home_mimo/modules/"

‎examples/mtp_trainer/test_dapo_mimo_7b_with_mtp_math_megatron.sh‎

Lines changed: 144 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,144 @@
1+
#!/usr/bin/env bash
2+
3+
set -xeuo pipefail
4+
5+
project_name='DAPO'
6+
exp_name='DAPO-mimo-7b-rl-megatron'
7+
8+
adv_estimator=grpo
9+
10+
use_kl_in_reward=False
11+
kl_coef=0.0
12+
use_kl_loss=False
13+
kl_loss_coef=0.0
14+
15+
clip_ratio_low=0.2
16+
clip_ratio_high=0.28
17+
18+
max_prompt_length=$((1024 * 2))
19+
max_response_length=$((1024 * 8))
20+
enable_overlong_buffer=True
21+
overlong_buffer_len=$((1024 * 4))
22+
overlong_penalty_factor=1.0
23+
24+
loss_agg_mode="token-mean"
25+
26+
train_prompt_bsz=128
27+
n_resp_per_prompt=16
28+
train_prompt_mini_bsz=32
29+
30+
# Ray
31+
# RAY_ADDRESS=${RAY_ADDRESS:-"http://localhost:8265"}
32+
# WORKING_DIR=${WORKING_DIR:-"${PWD}"}
33+
# RUNTIME_ENV=${RUNTIME_ENV:-"${WORKING_DIR}/examples/mtp_trainer/runtime_env.yaml"}
34+
NNODES=${NNODES:-16}
35+
NGPUS_PER_NODE=${NGPUS_PER_NODE:-8}
36+
# Paths
37+
RAY_DATA_HOME=${RAY_DATA_HOME:-"${HOME}/verl"}
38+
# very important! please modify the max_position_embeddings in config.json to 32768 after downloading from huggingface
39+
MODEL_PATH=${MODEL_PATH:-"${RAY_DATA_HOME}/models/MiMo-7B-RL"}
40+
CKPTS_DIR=${CKPTS_DIR:-"${RAY_DATA_HOME}/ckpts/${project_name}/${exp_name}"}
41+
TRAIN_FILE=${TRAIN_FILE:-"${RAY_DATA_HOME}/data/dapo-math-17k.parquet"}
42+
TEST_FILE=${TEST_FILE:-"${RAY_DATA_HOME}/data/aime-2024.parquet"}
43+
44+
# Algorithm
45+
temperature=1.0
46+
top_p=1.0
47+
top_k=-1 # 0 for HF rollout, -1 for vLLM rollout
48+
val_top_p=0.7
49+
50+
# Performance Related Parameter
51+
use_dynamic_bsz=True
52+
actor_ppo_max_token_len=$(((max_prompt_length + max_response_length) * 2))
53+
infer_ppo_max_token_len=$(((max_prompt_length + max_response_length) * 3))
54+
offload=True
55+
gen_tp=4
56+
train_tp=2
57+
train_pp=2
58+
train_cp=2
59+
60+
common_params=(
61+
actor_rollout_ref.model.mtp.enable=True
62+
actor_rollout_ref.model.mtp.enable_train=True
63+
actor_rollout_ref.model.mtp.mtp_loss_scaling_factor=0.1
64+
actor_rollout_ref.model.mtp.detach_encoder=True
65+
)
66+
67+
python -m verl.trainer.main_ppo \
68+
--config-path=config \
69+
--config-name='ppo_megatron_trainer.yaml' \
70+
data.train_files="${TRAIN_FILE}" \
71+
data.val_files="${TEST_FILE}" \
72+
data.prompt_key=prompt \
73+
data.truncation='left' \
74+
data.max_prompt_length=${max_prompt_length} \
75+
data.max_response_length=${max_response_length} \
76+
data.train_batch_size=${train_prompt_bsz} \
77+
actor_rollout_ref.rollout.n=${n_resp_per_prompt} \
78+
algorithm.adv_estimator=${adv_estimator} \
79+
algorithm.use_kl_in_reward=${use_kl_in_reward} \
80+
algorithm.kl_ctrl.kl_coef=${kl_coef} \
81+
actor_rollout_ref.actor.use_kl_loss=${use_kl_loss} \
82+
actor_rollout_ref.actor.kl_loss_coef=${kl_loss_coef} \
83+
actor_rollout_ref.actor.clip_ratio_low=${clip_ratio_low} \
84+
actor_rollout_ref.actor.clip_ratio_high=${clip_ratio_high} \
85+
actor_rollout_ref.actor.clip_ratio_c=10.0 \
86+
actor_rollout_ref.actor.ppo_micro_batch_size_per_gpu=2 \
87+
actor_rollout_ref.ref.log_prob_micro_batch_size_per_gpu=4 \
88+
actor_rollout_ref.rollout.log_prob_micro_batch_size_per_gpu=4 \
89+
actor_rollout_ref.model.path="${MODEL_PATH}" \
90+
actor_rollout_ref.actor.optim.lr=1e-6 \
91+
actor_rollout_ref.actor.optim.lr_warmup_steps=10 \
92+
actor_rollout_ref.actor.optim.weight_decay=0.1 \
93+
actor_rollout_ref.actor.ppo_mini_batch_size=${train_prompt_mini_bsz} \
94+
actor_rollout_ref.actor.megatron.param_offload=${offload} \
95+
actor_rollout_ref.actor.megatron.optimizer_offload=${offload} \
96+
actor_rollout_ref.actor.megatron.grad_offload=${offload} \
97+
actor_rollout_ref.actor.megatron.pipeline_model_parallel_size=${train_pp} \
98+
actor_rollout_ref.actor.megatron.tensor_model_parallel_size=${train_tp} \
99+
actor_rollout_ref.actor.megatron.context_parallel_size=${train_cp} \
100+
actor_rollout_ref.actor.entropy_coeff=0 \
101+
actor_rollout_ref.actor.optim.clip_grad=1.0 \
102+
actor_rollout_ref.actor.loss_agg_mode=${loss_agg_mode} \
103+
actor_rollout_ref.rollout.gpu_memory_utilization=0.80 \
104+
actor_rollout_ref.rollout.tensor_model_parallel_size=${gen_tp} \
105+
actor_rollout_ref.rollout.enable_chunked_prefill=True \
106+
actor_rollout_ref.rollout.max_num_batched_tokens=$((max_prompt_length + max_response_length)) \
107+
actor_rollout_ref.rollout.temperature=${temperature} \
108+
actor_rollout_ref.rollout.top_p=${top_p} \
109+
actor_rollout_ref.rollout.top_k=${top_k} \
110+
actor_rollout_ref.rollout.val_kwargs.temperature=${temperature} \
111+
actor_rollout_ref.rollout.val_kwargs.top_p=${val_top_p} \
112+
actor_rollout_ref.rollout.val_kwargs.top_k=${top_k} \
113+
actor_rollout_ref.rollout.val_kwargs.do_sample=True \
114+
actor_rollout_ref.rollout.val_kwargs.n=1 \
115+
actor_rollout_ref.rollout.name=sglang \
116+
actor_rollout_ref.ref.megatron.pipeline_model_parallel_size=${train_pp} \
117+
actor_rollout_ref.ref.megatron.tensor_model_parallel_size=${train_tp} \
118+
actor_rollout_ref.ref.megatron.context_parallel_size=${train_cp} \
119+
actor_rollout_ref.ref.megatron.param_offload=${offload} \
120+
reward_model.reward_manager=dapo \
121+
+reward_model.reward_kwargs.overlong_buffer_cfg.enable=${enable_overlong_buffer} \
122+
+reward_model.reward_kwargs.overlong_buffer_cfg.len=${overlong_buffer_len} \
123+
+reward_model.reward_kwargs.overlong_buffer_cfg.penalty_factor=${overlong_penalty_factor} \
124+
+reward_model.reward_kwargs.overlong_buffer_cfg.log=False \
125+
+reward_model.reward_kwargs.max_resp_len=${max_response_length} \
126+
trainer.logger='["console","tensorboard"]' \
127+
trainer.project_name="${project_name}" \
128+
trainer.experiment_name="${exp_name}" \
129+
trainer.n_gpus_per_node="${NGPUS_PER_NODE}" \
130+
trainer.nnodes="${NNODES}" \
131+
trainer.val_before_train=False \
132+
trainer.test_freq=10 \
133+
trainer.save_freq=-1 \
134+
trainer.total_epochs=10 \
135+
trainer.resume_mode=auto \
136+
trainer.log_val_generations=10 \
137+
actor_rollout_ref.rollout.disable_log_stats=False \
138+
actor_rollout_ref.rollout.prometheus.enable=True \
139+
actor_rollout_ref.rollout.prometheus.port=44398 \
140+
actor_rollout_ref.model.trust_remote_code=True \
141+
data.trust_remote_code=True \
142+
trainer.total_training_steps=400 \
143+
actor_rollout_ref.actor.megatron.use_mbridge=True \
144+
"${common_params[@]}"

‎verl/experimental/fully_async_policy/README.md‎

Lines changed: 3 additions & 3 deletions
Original file line numberDiff line numberDiff line change
@@ -65,7 +65,7 @@ The overall architecture of fully_async_policy is shown in the figure below. ful
6565
parts: Rollouter, MessageQueue, Trainer, and ParameterSynchronizer.
6666

6767
![fully_async_policy_structure](
68-
https://github.com/ArronHZG/verl-community/blob/recipe/async_policy/docs/fully_async_policy_structure.svg?raw=true)
68+
https://github.com/ArronHZG/verl-community/blob/main/docs/fully_async_policy_structure.svg?raw=true)
6969

7070
1. Rollouter generates sequences sample by sample and puts the generated samples into the MessageQueue, with the
7171
production speed controlled by freshness.
@@ -82,7 +82,7 @@ are used),
8282
but the overlap in their time consumption reduces the end-to-end time consumption.
8383

8484
![fully_async_policy_revenue](
85-
https://github.com/ArronHZG/verl-community/blob/recipe/async_policy/docs/fully_async_policy_revenue.svg?raw=true)
85+
https://github.com/ArronHZG/verl-community/blob/main/docs/fully_async_policy_revenue.svg?raw=true)
8686

8787
## Usage
8888

@@ -248,7 +248,7 @@ https://github.com/ArronHZG/verl-community/blob/recipe/async_policy/docs/fully_a
248248
3. As shown in figure d;
249249

250250
![fully_async_policy_mode](
251-
https://github.com/ArronHZG/verl-community/blob/recipe/async_policy/docs/fully_async_policy_mode.svg?raw=true)
251+
https://github.com/ArronHZG/verl-community/blob/main/docs/fully_async_policy_mode.svg?raw=true)
252252

253253
### Key Metrics
254254

‎verl/experimental/fully_async_policy/README_zh.md‎

Lines changed: 3 additions & 3 deletions
Original file line numberDiff line numberDiff line change
@@ -46,7 +46,7 @@ rollout的训练， 通过合理设置资源分配情况、参数同步频率等
4646
fully_async_policy的整体架构如下图所示，fully_async_policy主要由Rollouter、MessageQueue、Trainer、ParameterSynchronizer四部分组成。
4747

4848
![fully_async_policy_structure](
49-
https://github.com/ArronHZG/verl-community/blob/recipe/async_policy/docs/fully_async_policy_structure.svg?raw=true)
49+
https://github.com/ArronHZG/verl-community/blob/main/docs/fully_async_policy_structure.svg?raw=true)
5050

5151
1. Rollouter逐样本生成序列，并将生成的sample放入MessageQueue中，生产的速度受新鲜度控制。
5252
2. MessageQueue用于暂存Rollouter生成的sample。
@@ -59,7 +59,7 @@ https://github.com/ArronHZG/verl-community/blob/recipe/async_policy/docs/fully_a
5959
但是相互之间的耗时overlap，端到端的耗时反而有所缩减。
6060

6161
![fully_async_policy_revenue](
62-
https://github.com/ArronHZG/verl-community/blob/recipe/async_policy/docs/fully_async_policy_revenue.svg?raw=true)
62+
https://github.com/ArronHZG/verl-community/blob/main/docs/fully_async_policy_revenue.svg?raw=true)
6363

6464
## 使用方式
6565

@@ -199,7 +199,7 @@ https://github.com/ArronHZG/verl-community/blob/recipe/async_policy/docs/fully_a
199199
3. 如图d所示；
200200

201201
![fully_async_policy_mode](
202-
https://github.com/ArronHZG/verl-community/blob/recipe/async_policy/docs/fully_async_policy_mode.svg?raw=true)
202+
https://github.com/ArronHZG/verl-community/blob/main/docs/fully_async_policy_mode.svg?raw=true)
203203

204204
### 关键指标
205205

‎verl/models/mcore/model_forward.py‎

Lines changed: 13 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -17,6 +17,7 @@
1717
import torch
1818

1919
from verl.utils.megatron_utils import unwrap_model
20+
from verl.workers.config import MtpConfig
2021

2122
from .util import (
2223
postprocess_bshd,
@@ -41,6 +42,7 @@ def model_forward(
4142
logits_processor_args: dict = None,
4243
value_model=False,
4344
data_format: str = "thd",
45+
mtp_config: MtpConfig = None,
4446
):
4547
"""Forward pass for models with sequence packing."""
4648
assert data_format in ["thd", "bshd"], "data_format must be 'thd' or 'bshd'"
@@ -65,10 +67,19 @@ def model_forward(
6567
batch_size, seq_len = attention_mask.shape[:2]
6668
if data_format == "thd":
6769
input_ids_rmpad, packed_seq_params = preprocess_packed_seqs(
68-
input_ids, attention_mask, pre_process=pre_process, use_fp8_padding=use_fp8_padding
70+
input_ids, attention_mask, pre_process=pre_process or post_process, use_fp8_padding=use_fp8_padding
6971
)
7072
input_ids_rmpad = input_ids_rmpad.contiguous()
7173

74+
# when pp > 1 and processor is not None, we need to pass the labels and loss_mask to the model
75+
if mtp_config and mtp_config.enable_train and post_process:
76+
args = {
77+
k: preprocess_packed_seqs(v, attention_mask, pre_process=True, use_fp8_padding=use_fp8_padding)[0]
78+
for k, v in logits_processor_args.items()
79+
}
80+
model_kwargs["labels"] = args["label"].contiguous()
81+
model_kwargs["loss_mask"] = args["label_mask"].contiguous()
82+
7283
input_args = dict(
7384
input_ids=input_ids_rmpad,
7485
attention_mask=None,
@@ -86,6 +97,7 @@ def model_forward(
8697
input_args["attention_mask"] = attention_mask
8798

8899
output_orig = model(**input_args)
100+
89101
if post_process and logits_processor is not None:
90102
args = {
91103
k: preprocess_packed_seqs(v, attention_mask, pre_process=True, use_fp8_padding=use_fp8_padding)[0]

0 commit comments

Comments
 (0)