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

Convert a DataProtoItem to a DataProto object

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]

Fetch JSON data from a URL with retries.

Calculate the percentile of a given rating.

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.

Return the result of gcd() or lcm(), as UnevaluatedExpr f: str - name of function ("gcd" or "lcm") args: List[Expr] - list of function a