Pop a subset of the DataProto via `batch_keys` and `meta_info_keys` Args: batch_keys (list, optional): a list of strings indicati

move the batch to device Args: device (torch.device, str): torch device Returns: DataProto: the current Data

Update kl_coef according to current KL.

r""" Instantiate a [`GenerationConfig`] from a generation configuration file. Args: pretrained_model_name (`str` or `os.P

Report the current GPU VRAM usage.

Get ulysses sequence parallel process group.

Instantiates a [`GenerationConfig`] from a Python dictionary of parameters. Args: config_dict (`Dict[str, Any]`):

Split the batch among dim=0 into chunks. The meta_info is passed to each DataProto after split. Args: chunks (int): the number of

r""" Save a generation configuration object to the directory `save_directory`, so that it can be re-loaded using the [`~GenerationConf

Union with another DataProto. Union batch and meta_info separately. Throw an error if - there are conflict keys in batch and they are

Concat a list of DataProto. The batch is concatenated among dim=0. The meta_info is assumed to be identical and will use the first one.

Create a DataProto from a dict of tensors. This assumes that 1. All the tensor in tensors have the same dim0 2. Only dim0 is the batch

Downloads a file from the Hugging Face Hub and returns its contents as a dictionary. Parameters: - file_name (str): The name of the

Select a subset of the DataProto via batch_keys and meta_info_keys Args: batch_keys (list, optional): a list of strings indicatin

Serializes this instance to a Python dictionary.

Serializes this instance to a JSON string. Args: use_diff (`bool`, *optional*, defaults to `True`): If s

Append dict to a dict of list.

Get the resource pool of the worker.

Performs a single optimization step. Args: closure (callable, optional): A closure that reevaluates the model and return

A func to sync embedding input with alltoall in sequence parallel gather sequence dimension and scatter head dim: e.g. seq_dim: 1, head_d

Create a huggingface pretrained processor.

Get ulysses sequence parallel world size.

Compute mean of tensor with a masked values.

Pad a nD tensors in the last dim to max_seq_len.

Pad and slice input_ids to be divisible by sp_size Pad position_ids to be divisible by sp_size. Note both input_ids_rmpad and position_i

Updates attributes of this class instance with attributes from `kwargs` if they match existing attributes, returning all the unused k

Validates the values of the attributes of the [`GenerationConfig`] instance. Raises exceptions in the presence of parameterization th

Returns: log_probs: # (bs, response_len)

Compute logits given a batch of data. Args: data (DataProto): a batch of data represented by DataProto. It must contain key ```in

Checks whether the passed dictionary and its nested dicts have a *torch_dtype* key and if it's not None, converts torch.dtype to a st

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

Instantiates a [`GenerationConfig`] from a [`PretrainedConfig`]. This function is useful to convert legacy [`PretrainedConfig`] objec

Get the number of gpus in this cluster.

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

Create a huggingface pretrained tokenizer.

Wraps the mathruler 'grade_answer' function with a timeout. If the function takes too long, it returns 'TIMED_OUT' instead of hanging.

This wrapper applies a timeout to each individual `grade_answer` call. If the function's execution exceeds the specified timeout, it will ret

Repeat the batch data a specified number of times. Args: repeat_times (int): Number of times to repeat the data.

Set ulysses sequence parallel process group.

Serializes this instance to a Python dictionary. Returns: `Dict[str, Any]`: Dictionary of all the attributes that make u

Reorder the data on single controller such that each dp rank gets similar total tokens

Bind the worker method to the WorkerGroup

Binds the methods of each worker to the WorkerDict. Note that we only bind public methods that are decorated by register