Build the memory buffer given weight_buffer_meta Args: weight_buffer_meta: contains mapping from name to a dictionary containing shape an

for cuda memory alignment, make sure alignment by 128-bits

The training loop of PPO. The driver process only need to call the compute functions of the worker group through RPC to construct the

Init resource pool and worker group

Pad a DataProto to size divisible by size_divisor Args: size_divisor (int): size divisor Returns: data: (DataProto): the pad

Note that this operation is in-place

Set input tensor to be used instead of forward()'s input. When doing pipeline parallelism the input from the previous stage comes fro

Union two tensordicts.

TODO(zhangchi.usc1992): We may change the implementation later

Returns: entropy: # (bs, response_len) log_probs: # (bs, response_len)

allgather params of all pp ranks. Return a list of handles

This function should return a class instance that delegates the calls to every cls in cls_dict

Converts a string to a sequence of ids (integer), using the tokenizer and vocabulary. Args: text (`str`, `List[str]` or

Estimate the FLOPS based on the number of valid tokens in the current batch and the time taken. Args: batch_seqlens (Lis

Args: input_ids: input ids. shape (batch_size, seq_length) attention_mask: attention_mask. shape (batch_size, seq_le

We assume: - The model takes input: (input_ids, attention_mask, position_ids). No rmpad for the input - The communication sha

A func to sync attention result with alltoall in sequence parallel gather head dimension and scatter seq dim: e.g. seq_dim: 1, head_dim:

Get all the parameters of the models in all pp ranks Returns: params: List[List[Dict[str, Tensor]]]: a list of parameters in all

Create a schedule with a learning rate that decreases following the values of the cosine function between the initial lr set in the optimizer

e.g. end of sentence token=1 response_id: [0, 0, 2, 42, 3, 5, 1, 0, 0] eos_mask: [1, 1, 1, 1, 1, 1, 1, 0, 0]

get order of seq lengths to make partitions balanced, this is used in balacing sum of seqlength across dp ranks and microbatches Paramete

Return a dictionary containing name to a shape and dtype.

load all model params to cuda

Compute the log_probs from logits with rmpad input_ids and logits. Note that logits_rmpad = model(input_ids_rmpad). For each sentences, there is a

Make minibatch iterator for updating the actor Args: data (DataProto): a DataProto containing keys ``input_ids``

Normalize the pp vpp params into a complete named parameters. This is useful when gather parameters from pp ranks and passed to a model with

AllGather data from sp region This is because the data is first sharded along the FSDP dimension as we utilize the DP_COMPUTE