output = model(input_ids=input_ids, attention_mask=None, position_ids=position_ids)

def preprocess_packed_seqs(

    input_ids: torch.Tensor,

    labels: torch.Tensor,

    attention_mask_1d: torch.Tensor,

    tp_size: int

) -> Tuple[torch.Tensor, torch.Tensor, torch.Tensor, torch.Tensor]:

    """

    处理逻辑：

    ● 输入的input_ids，忽略原来的padding，根据tp_size先padding，使每个seq能被tp_size整除，然后把所有padding后的seq拼接到一起。label处理和input_ids对应（类似）。

    ● 记录position_ids_packed用于后面做模型输入。

    ● 记录seqlens_in_batch、cu_seqlens_padded用于后面恢复原始input_ids的格式。

    输入：

    ● input_ids：tensor,shape为(batch_size, seq_len),存放的是token IDs

    ● attention_mask_1d：tensor,shape为 (batch_size, seq_len)，其中的1表示token，0表示padding

    ● tp_size:  为了使seq_len能被tp_size整除，会做padding

    输出：可参照上图

    ● input_ids_packed： shape为(1, pack_length)，处理后的新input_ids。

    ● position_ids_packed：会作为模型的输入(对应参数position_ids，用于生成postion_embedding)。shape为(1, pack_length)，记录每个padding后的seq的位置。

    ● seqlens_in_batch：shape为 (batch_size,)，记录每个sequence的原始seq_len（不算padding），用于后面恢复原始输入shape

    ● cu_seqlens_padded：shape为 (batch_size+1,)  记录每个子padding后的seq的起止位置，用于后面恢复原始输入shape。

    """

def postprocess_packed_seqs(

    output: torch.Tensor,

    seqlens_in_batch: torch.Tensor,

    cu_seqlens_padded: torch.Tensor,

    seq_len: int,

    prompt_length: torch.Tensor = None

) -> torch.Tensor:

    """

    输入：

    ● output： shape为(1, pack_length, ...)，preprocess_packed_seqs处理后的input_ids_packed经过模型后的输出

    ● seqlens_in_batch：shape为 (batch_size,)，preprocess_packed_seqs计算结果

    ● cu_seqlens_padded：shape为 (batch_size+1,)，preprocess_packed_seqs计算结果

    ● seq_len：int类型，preprocess_packed_seqs处理之前的input_ids的seq_len（最长的seq的长度，包括padding），用于恢复preprocess_packed_seqs处理前的shape

    输出：

    ● output_new：shape为(batch_size, seq_len, ...)，为恢复后的输出

    """

def rearrange_micro_batches(seqlen_list: List[int], max_token_len: int, dynamic_max_batch_size=None, dp_group=None):

# Calculate the minimum number of bins

total_sum_of_seqlen = sum(seqlen_list)

k_partitions = (total_sum_of_seqlen + max_token_len - 1) // max_token_len

#先计算batch中所有seq的长度和，再除max_token_len，得到要分的micro_batch的数量，每个macro_batch使用remove_padding拼接序列后，长度不会超过限制。k_partitions就是分几份

if dynamic_max_batch_size is not None:

    k_partitions = max(k_partitions, (len(seqlen_list) + dynamic_max_batch_size - 1) // dynamic_max_batch_size)

def rearrange_micro_batches(seqlen_list: List[int], max_token_len: int, dynamic_max_batch_size=None, dp_group=None):

    ......

    # 根据总序列长度和指定的max_token_len，计算最少分几组，每组长度不超过max_token_len

    total_sum_of_seqlen = sum(seqlen_list)

    k_partitions = (total_sum_of_seqlen + max_token_len - 1) // max_token_len

    # 根据dynamic_max_batch_size的限制，再算可以分几组，然后综合上面的结果取最大值

    if dynamic_max_batch_size is not None:

        k_partitions = max(k_partitions, (len(seqlen_list) + dynamic_max_batch_size - 1) // dynamic_max_batch_size)

    #如果指定dp_group，可以在dp_group内同步k_partitions，保证dp_group内划分分组数一致    

    if dist.is_initialized():

        k_partitions = torch.tensor([k_partitions], device='npu')

        dist.all_reduce(k_partitions, op=dist.ReduceOp.MAX, group=dp_group)

        k_partitions = k_partitions.cpu().item()

    #使用karmarkar_karp算法来划分分组

    partitions = karmarkar_karp(seqlen_list=seqlen_list, k_partitions=k_partitions, equal_size=False)

    #返回indexes列表，其中存的是每个partition(micro_batch)在原始序列中的index，用于后面恢复

    return partitions

def _forward_backward_batch(self, batch: Dict[str, torch.Tensor], forward_only: bool = False):

    #按照动态批大小算法切分batch

    if self.use_dynamic_bsz:

        batches, indices = self._split_batches_with_dynamic_bsz(batch, self.max_packing_token_size, self.dynamic_max_batch_size)

    #普通batch切分方法，按照micro_batch_size切分

    else:

        batches = self._split_batches(batch, batch_size=self.micro_batch_size,

                                        shuffle_mini_batch=self.shuffle_mini_batch)

    ......

    data_iter = iter(batches)

    ......

    # batch should be a list of batches inside micro-batches

    #数据前后向处理

    losses_reduced = self.forward_backward_func(

        forward_step_func=forward_step,

        data_iterator=data_iter,

        model=self.model,

        num_microbatches=n_micro_batch,

        seq_length=self.micro_batch_size * seq_len if self.use_remove_padding else seq_len,

        micro_batch_size=1 if self.use_remove_padding else self.micro_batch_size,

        forward_only=forward_only,

        collect_non_loss_data=forward_only,

    )

    #动态批大小算法反向操作，恢复数据原始分组

    # Reverse the batch index to be the same outside

    if self.use_dynamic_bsz and forward_only and post_process:

        losses_reduced_list = torch.cat(losses_reduced, dim=0)

        indices = list(itertools.chain.from_iterable(indices))

        revert_indices = get_reverse_idx(indices)

        losses_reduced = [losses_reduced_list[[idx, ]] for idx in revert_indices]

    return losses_reduced

def forward_step(batch_iter, model):

    .......

    if self.use_remove_padding:

        #remove padding前处理

        input_ids, position_ids, process_batch, seqlens_in_batch, cu_seqlens_padded, index = self._get_forward_batch_info(batch_iter)

        #数据输入模型处理

        output = model(input_ids=input_ids, attention_mask=None, position_ids=position_ids)

        output.div_(self.temperature)

        #模型输出的使用（包括后处理）

        return output, partial(self.loss_func.compute_loss,

                                batch=process_batch,

                                forward_only=forward_only,

                                use_remove_padding=self.use_remove_padding,

                                seqlens_in_batch=seqlens_in_batch,

                                cu_seqlens_padded=cu_seqlens_padded,

                                seq_len=seq_len,

                                use_dynamic_bsz=self.use_dynamic_bsz,

                                actual_micro_batch_size=batch_size / n_micro_batch,

                                index=index)

    ......

#mindspeed_rl.models.base.base_training_engine.BaseTrainingEngine._forward_backward_batch    

if self.use_dynamic_bsz:

    batches, indices = self._split_batches_with_dynamic_bsz(batch, self.max_packing_token_size, self.dynamic_max_batch_size)

else:

    batches = self._split_batches(batch, batch_size=self.micro_batch_size,

                                    shuffle_mini_batch=self.shuffle_mini_batch)

#mindspeed_rl.models.base.base_training_engine.BaseTrainingEngine._split_batches_with_dynamic_bsz

def _split_batches_with_dynamic_bsz(batch: Dict, max_packing_token: int, dynamic_max_batch_size: int) -> tuple[List[Dict], List[List[int]]]:

    seq_len_list = []

    #计算batch中数据的prompt_length+response_length生成seq_len_list

    for prompt_len, response_len in zip(batch['prompt_length'], batch['response_length']):

        seq_len_list.append(prompt_len.item() + response_len.item())

    #调用rearrange_micro_batches获取冲切分后的分组

    partitions = rearrange_micro_batches(seq_len_list, max_packing_token, dynamic_max_batch_size=dynamic_max_batch_size)

    #将batch数据按照新分组方式重组

    batches = []

    for key, tensors in batch.items():

        for batch_idx, partition in enumerate(partitions):

            if batch_idx >= len(batches):

                batches.append({})

            batches[batch_idx][key] = tensors[partition]

    #返回新分组、以及新分组数据对应原始序列中的index，用于后面恢复

    return batches, partitions

#mindspeed_rl.models.base.base_training_engine.BaseTrainingEngine._forward_backward_batch        

def _forward_backward_batch(...):

    ...

    elif self.use_remove_padding:

        input_ids, position_ids, process_batch, seqlens_in_batch, cu_seqlens_padded, index = self._get_forward_batch_info(batch_iter)  

    ...

#mindspeed_rl.models.base.base_training_engine.BaseTrainingEngine._get_forward_batch_info

def _get_forward_batch_info(self, batch_iter):

    batch = next(batch_iter)

    input_ids = batch['input_ids']

    ......

    attention_mask_1d = generate_mask(input_ids, batch['prompt_length'] + batch['response_length']).to(input_ids.device)

    ......

    if self.use_remove_padding:

        tp_size = get_parallel_state().get_tensor_model_parallel_world_size()

        if self.megatron_config.context_parallel_algo == "megatron_cp_algo":

            multi = 2 * tp_size * cp_size

        else:

            multi = tp_size * cp_size

        input_ids, position_ids, labels, seqlens_in_batch, cu_seqlens_padded = preprocess_packed_seqs(

            input_ids=input_ids, labels=labels, attention_mask_1d=attention_mask_1d, tp_size=multi)         

#mindspeed_rl.models.base.base_training_engine.BaseTrainingEngine._forward_backward_batch  

losses_reduced = self.forward_backward_func(

    forward_step_func=forward_step,

    data_iterator=data_iter,

    model=self.model,

    num_microbatches=n_micro_batch,

    seq_length=self.micro_batch_size * seq_len if self.use_remove_padding else seq_len,

    micro_batch_size=1 if self.use_remove_padding else self.micro_batch_size,

    forward_only=forward_only,

    collect_non_loss_data=forward_only,

)

#mindspeed_rl.models.base.base_training_engine.BaseTrainingEngine._forward_backward_batch        

def forward_step(batch_iter, model):

    .......

    if self.use_remove_padding:

        #remove padding前处理

        input_ids, position_ids, process_batch, seqlens_in_batch, cu_seqlens_padded, index = self._get_forward_batch_info(batch_iter)

        #数据输入模型处理

        output = model(input_ids=input_ids, attention_mask=None, position_ids=position_ids)

        output.div_(self.temperature)

        #模型输出的使用（包括后处理）

        return output, partial(self.loss_func.compute_loss,

                                batch=process_batch,

                                forward_only=forward_only,

                                use_remove_padding=self.use_remove_padding,

                                seqlens_in_batch=seqlens_in_batch,

                                cu_seqlens_padded=cu_seqlens_padded,

                                seq_len=seq_len,

                                use_dynamic_bsz=self.use_dynamic_bsz,

                                actual_micro_batch_size=batch_size / n_micro_batch,

                                index=index)

    ......

#mindspeed_rl.models.loss.base_loss_func.BaseLossFunc.compute_log_probs

log_probs = compute_log_probs(output, labels)

......

if use_remove_padding:

    log_probs_allgather = get_tensor_allgather_cp_with_pack(log_probs, cp_size, index)

    seqlens_in_batch = kwargs.get('seqlens_in_batch', None)

    cu_seqlens_padded = kwargs.get('cu_seqlens_padded', None)

    seq_len = batch['responses'].shape[-1]

    log_probs = postprocess_packed_seqs(log_probs_allgather, seqlens_in_batch, cu_seqlens_padded, seq_len, prompt_length=batch['prompt_length'])

    if not skip_entropy:

        entropy = vocab_parallel_entropy(output)

        entropy = postprocess_packed_seqs(entropy, seqlens_in_batch, cu_seqlens_padded, seq_len, prompt_length=batch['prompt_length'])

    else:

        entropy = torch.zeros_like(log_probs)

    return log_probs, entropy

#mindspeed_rl.models.loss.grpo_actor_loss_func.GRPOActorLossFunc.compute_loss

# compute policy loss

pg_loss, pg_clipfrac, ppo_kl, kl_loss, entropy_loss = self._compute_grpo_policy_loss(old_log_prob=old_log_prob,

                             log_prob=log_probs,

                             ref_log_prob=ref_log_prob,

                             advantages=advantages,

                             entropy=entropy,

                             eos_mask=response_mask,

                             cliprange=self.clip_ratio,

                             kl_ctrl=self.kl_ctrl,

                             kl_penalty=self.kl_penalty,

                             entropy_coeff=self.entropy_coeff)

use_dynamic_bsz = kwargs.get('use_dynamic_bsz', False)

actual_micro_batch_size = kwargs.get('actual_micro_batch_size', None)

if use_dynamic_bsz and not forward_only:

    policy_loss = pg_loss * (batch['responses'].size(0) / actual_micro_batch_size)

else:

    policy_loss = pg_loss

#mindspeed_rl.models.base.base_training_engine.BaseTrainingEngine._forward_backward_batch

# Reverse the batch index to be the same outside

if self.use_dynamic_bsz and forward_only and post_process:

    losses_reduced_list = torch.cat(losses_reduced, dim=0)

    indices = list(itertools.chain.from_iterable(indices))

    revert_indices = get_reverse_idx(indices)

    losses_reduced = [losses_reduced_list[[idx, ]] for idx in revert_indices]