Check if the given software and CANN versions support IPC. Compares the software version and CANN toolkit version against minimum required ve

Get all parameters, deduplicated by object id to avoid double-optimizing tied weights.

Sleep all rollout replicas: free weight and kv_cache device memory.

Check correctness of code generation with a global timeout. The global timeout is to catch some extreme/rare cases not handled by the timeouts

Get the distributed communication backend based on device type. Returns the appropriate collective communication backend for the detected acc

Slice response from unpad model output. Args: tensor: a nested tensor or a tensor of shape (total_nnz,*), total_nnz is the to

In megatron, the model and optimizer storage are: - bf16 parameter data chunked in model parallel group - fp32 grad chunked in model para

Create mock rollout and actor engine. Args: peft_config: If not None, get_per_tensor_param returns this as peft_config. Use a

Reproduce the update_weights() logic from engine_workers.py. This mirrors the actual code so we can test the branching without importing the

Get ulysses sequence parallel process group.

Append values from new_data to lists in data. For each key in new_data, this function appends the corresponding value to a list stored under

Convert a config dictionary to CLI arguments for vLLM server. Handles different value types appropriately: - None: skipped - bool Tr

Get ulysses sequence parallel world size.

Initialize class state shared across all instances.

Load megatron model to GPU. Args: models: The model to load. load_grad: Whether to load gradients. load_frozen_params

Unified entry point for `unpad_input` across CUDA and NPU backends. Dynamically dispatches to the appropriate device-specific implementation

Make a HTTP request with retry logic and consistent error handling. Args: endpoint (str): The API endpoint to call (without leadi

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

Load an object from a module path. Args: module_path (str): See :func:`load_module`. object_name (str): The name of t

Preprocess packed sequences CP splits sequence into CP*2 chunks, and each GPU gets 2 chunks (GPU0 gets first and last chunks, GPU1 gets s

Get CPU tensor.

Helper to create DataProto for testing.

Create a DataProto from a dict of tensors and non_tensors

Generates a response for the current turn of interaction. Returns a tuple containing: - should_terminate_sequence (bool): Tru

Convert TensorDict from left-right padding to no-padding format. Args: data: TensorDict with "input_ids", "attention_mask", "respons

Function to create a new training engine instance based on the provided config. Args: key: A configuration object contain

Unified entry point for `rearrange` across CUDA and NPU backends. Dynamically dispatches to the appropriate device-specific implementation:

Determine if checkpoint should be saved based on capacity esi expiration. Args: max_steps_duration: Max estimated time (seconds) req

Split a RayResourcePool into multiple SubRayResourcePool. resouce_pool can also be a SubRayResourcePool (have been splited) for multiple-time

Make an async HTTP request with retry logic and consistent error handling. Args: endpoint (str): The API endpoint to call (withou

Execute a method on all workers synchronously. Args: method_name: Name of the method to execute *args: Positional arg

Find a free port on the given address. By default the socket is closed internally, suitable for immediate use. Set with_alive_sock=True to ke

Unified entry point for `index_first_axis` across CUDA and NPU backends. Dynamically dispatches to the appropriate device-specific implement

_summary_ Thread-safe directory creation function that ensures the directory is created even if multiple processes attempt to create it simult

Compute the mean of `values` over elements selected by `mask`. Args: values (Tensor): Input tensor. mask (Tensor): Boolean o

Given the config, do we need ref policy.

Prepare micro batches from data.

Set ulysses sequence parallel process group.

Start profiling on the replica.

Stop profiling on the replica.

Context manager for temporarily setting an environment variable. This context manager ensures that environment variables are properly set and res

Copied from ManiSkill! Maps any given sequence to a torch tensor on the CPU/GPU. If physx gpu is not enabled then we use CPU, otherwise G

Create a mock checkpoint directory.

Spawn sender + receiver processes, then validate received tensors.

Create a random mask given input_ids. Support left padding and right padding. Process: - Sample valid token length - Sample left_padding l

Create a dataset. Arguments: data_paths: List of paths to data files. data_config: The data config. tokenizer (Tokenizer)

Create shared memory for weight transfer. If already exists, attach to it.

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

Record the current training step based on the new batch. Args: new_batch (DataProto): The new batch of data being processed.

Reduces a dictionary of metric lists by computing the mean, max, or min of each list. The reduce operation is determined by the key name:

Move model and/or optimizer to CPU or GPU.

Update weights from trainer to rollout replicas. Args: global_steps: The global steps of the trainer.

model: AutoModelForCausalLM

Create a sampler for the dataset. Arguments: data_config: The data config. dataset (Dataset): The dataset. Returns:

Gather a tensor across a process group and optionally unpad its padded elements. Args: x (Tensor): Input tensor to gather. g

Init standalone rollout server, create new resource pool for this rollout.

Whether the current rank is the first rank in model parallel group that contains model outputs

Release GPU memory occupation temporarily. Args: tags (Optional[List[str]], optional): List of tags to specify which memory to re

Release or rebuild GPU buffers for MegaTrain. MegaTrain keeps all parameters on CPU. The GPU only holds transient compute buffers (do

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

broadcast tensor from rank0 across mp_group

Calculate global rank with support for CP/EP parallelism