Query the dispatch info for a given mesh name. Args: mesh_name (str): Name of the mesh to query dispatch info for

Inner function that handles the core timing logic. Args: name (str): The name/identifier for this timing measurement. timing_raw

Update the input_ids, attention_mask, position_ids, and loss_mask of the request in additive manner.

Acquire a token from the bucket.

添加过滤函数 Args: filter_func: 过滤函数，接收一个字典，返回布尔值 返回 True 表示保留该数据，False 表示过滤掉 name

Apply KL penalty to the token-level rewards. This function computes the KL divergence between the reference policy and current policy, then a

Performs bootstrap resampling to estimate statistics of metrics. This function uses bootstrap resampling to estimate the mean and standard d

Calculate approximate computational workload for transformer attention. Estimates FLOPs for dense transformer blocks based on sequence length usi

生成一个迷宫谜题 参数: size: 迷宫大小，(行数, 列数) start_pos: 起点位置 end_pos: 终点位置，如果为None则默认为右下角 diffic

Compute advantage estimates for policy optimization. This function computes advantage estimates using various estimators like GAE, GRPO, REINFORC

Compute the log probability of the responses given input_ids, attention_mask and position_ids Args: data (DataProto): a DataProto

Compute the clipped policy objective and related metrics for PPO. Adapted from https://github.com/huggingface/trl/blob/main/trl/trainer/

Compute the attention mask for the response part of the sequence. This function extracts the portion of the attention mask that corresponds to th

Create a dataset.

Compute the score for a given solution based on the data source. Args: data_source (str): The source dataset identifier which determines

Enable activation offloading for the model. It groups activations by TransformerLayer and offloads activation groups asynchronously. This mea

torch.nn.utils.clip_grad_norm_ cann't run on cpu parameter DTensor

Generate an n×m Campsite puzzle with at least one valid solution. 1) Randomly place k tents in the grid with no two tents adjacent o

Split input batch and dispatch to agent loop workers. Args: prompts (BatchMeta): Input batch. Returns: Batch

生成星战谜题的测试用例网格。 参数: size -- 网格的大小 (假设是正方形网格，size 表示边长) num_regions -- 区域的数量（等于星星数量） 返回: grid -- 生成的网格，每个区域由不同的数字表示，0

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

Calculate the pipeline sharding configuration for Megatron-LM. Args: layer_num: Total number of layers in the model. pp_size: Num

Partition samples into k groups while keeping samples with the same uid together. Args: seqlen_list: List of sequence lengths for ea

Get models according to different types of inference_engine Args: rollout: rollout object Returns: mo

Determine the number of transformer layers to build for the current pipeline stage. Args: config (TransformerConfig): Configuration o

Get the rollout class by name. Args: rollout_name: The name of the rollout. mode: The mode of the rollout, async: server mode.

Compute per-group mean/std/count in pure PyTorch. mean_g = sum / count std_g = sqrt( max( (sum2 - sum^2/count) / max(count-1, 1), eps )

Initialize tools from config file. Supports both NATIVE and MCP tool types. For MCP tools, a temporary event loop is created only when needed

Check if Ascend NPU is available for PyTorch operations. Attempts to detect NPU availability by checking for the torch.npu module and its is_

Launch http server in each node.

Load weights for mcore GPT model.

Given a role worker mapping, do we need reward model.