Load a function from a module. :param path: The path to the function, e.g. "module.submodule.function". :return: The function object.

Every metrics-dict value is a list of 1D tensor, i.e., [torch.Tensor] with shapes exactly the same as log_probs. All metrics will be aggreg

Wrap communication functions with memory check.

Reset the default value of a Megatron argument. :param parser: The argument parser. :param name: The name of the argument to reset. :

Runs the agent by sending a prompt to the LLM.

Get the Gloo group for distributed communication.

Log data to tensorboard Args: data (dict): Dictionary containing metric names and values step (int): Current step/epo

Resolve per-group defaults from model-level then args-level values.

Make a POST request to the specified endpoint with the given payload. Args: endpoint: The API endpoint to call payloa

Gather tensors across all ranks in the context parallel group. The first dimension of the output tensor will be the `response_length`.

Override to handle expert bias names correctly. Base class extracts expert_id by splitting on 'weight', which fails for bias paramete

All offsets start from the begining of the prompt.

Do one rollout training

Wrap query functions (get_rank, get_world_size, etc.) without memory check.

Build the final interaction result with all collected data. Args: res: InteractionResult instance to populate

Convert MTP layer parameters from MCore to HF format.

Return the router URL for a named model. Use this in custom rollout functions to route requests to a specific model when multiple models are

Return num_samples samples

Fire release_memory_occupation on all engines (non-blocking). Returns a list of Ray ObjectRefs. Skipped for groups that do not overl

Can be used either as a decorator or a context manager: @timer def func(): ... or with timer("block_name"): ..

Forward non-communication calls without memory check.

Map MCore weight names to Hugging Face weight names. Args: mcore_weights_name: MCore weight name Returns:

Load an image from path, ensure RGB, encode as PNG base64 string.

Calculate correct sample mean for CP