Repository: volcengine/verl
File: verl/models/mcore/model_forward_1f1b_overlap.py (Python)

Code before:

from megatron.core.models.common.model_chunk_schedule_plan import TransformerModelChunkSchedulePlan
from megatron.core.models.gpt.gpt_model import GPTModel
from megatron.core.utils import make_viewless_tensor
from torch import Tensor

from verl.models.mcore.util import preprocess_packed_seqs
from verl.utils.kernel.linear_cross_entropy import linear_cross_entropy
from verl.utils.megatron_utils import unwrap_model
from verl.utils.model import CausalLMOutputForPPO

from .util import postprocess_packed_seqs, postprocess_packed_seqs_for_dict_output


def gptmodel_forward_1f1b_overlap(
    model: GPTModel,
    input_ids: Tensor,
    position_ids: Tensor,
    attention_mask: Tensor,
    labels: Tensor = None,
    labels_mask: Tensor = None,
    sequence_parallel: bool = False,
    multi_modal_inputs: Optional[dict] = None,
    logits_processor: Optional[Callable] = None,
    logits_processor_args: Optional[dict] = None,
    temperature: float = 1.0,
) -> TransformerModelChunkSchedulePlan:
    pre_process: bool = unwrap_model(model).pre_process
    post_process: bool = unwrap_model(model).post_process
    assert logits_processor is None, "only support fused kernel"
    batch_size, seq_len = attention_mask.shape[:2]
    input_ids_rmpad, packed_seq_params = preprocess_packed_seqs(input_ids, attention_mask, pre_process=pre_process)
    input_ids_rmpad = input_ids_rmpad.contiguous()

    schedule_plan = model.build_schedule_plan(
        input_ids=input_ids_rmpad,
        attention_mask=attention_mask,
        labels=labels,
        position_ids=position_ids,
        packed_seq_params=packed_seq_params,
    )
    if post_process:
        attention_mask_out = attention_mask

        def _postprocess(
            self,
            hidden_states,
            input_ids,
            position_ids,
            labels,
            rotary_pos_emb,

Review comment (suggestion):
Similar to #3849, I think it would be better to have an additional `**kwargs` to support deepstack for Qwen3VL MoE models / any possible additional parameters for the future, and pass them to the decoder.

Code after:

import torch
from megatron.core.models.common.model_chunk_schedule_plan import TransformerModelChunkSchedulePlan
from megatron.core.models.gpt.gpt_model import GPTModel
from megatron.core.utils import make_viewless_tensor
from torch import Tensor

from verl.models.mcore.util import preprocess_packed_seqs
from verl.utils.kernel.linear_cross_entropy import linear_cross_entropy
from verl.utils.megatron_utils import unwrap_model
from verl.utils.model import CausalLMOutputForPPO

from .util import postprocess_packed_seqs, postprocess_packed_seqs_for_dict_output


def gptmodel_forward_1f1b_overlap(
    model: GPTModel,
    input_ids: Tensor,
    position_ids: Tensor,
    attention_mask: Tensor,
    labels: Tensor = None,
    labels_mask: Tensor = None,
    multi_modal_inputs: Optional[dict] = None,
    logits_processor: Optional[Callable] = None,
    logits_processor_args: Optional[dict] = None,
    temperature: float = 1.0,
) -> TransformerModelChunkSchedulePlan:
    pre_process: bool = unwrap_model(model).pre_process
    post_process: bool = unwrap_model(model).post_process
    assert logits_processor is None, "only support fused kernel"
    batch_size, seq_len = attention_mask.shape[:2]
    input_ids_rmpad, packed_seq_params = preprocess_packed_seqs(input_ids, attention_mask, pre_process=pre_process)
    input_ids_rmpad = input_ids_rmpad.contiguous()

    schedule_plan = model.build_schedule_plan(
        input_ids=input_ids_rmpad,
        attention_mask=attention_mask,
        labels=labels,
        position_ids=position_ids,
        packed_seq_params=packed_seq_params,
    )
    if post_process:
        attention_mask_out = attention_mask

        def _postprocess(
            self,
            hidden_states,
            input_ids,
            position_ids,
            labels,
            rotary_pos_emb,