Compute entropy when the logits are sharded in tp ranks Args: vocab_parallel_logits: (total_nnz, vocab_size // tp_size) Returns: (to

TODO(zhangchi.usc1992): We may change the implementation later

Initialize the VeOmniEngine. Sets up distributed device meshes, LoRA, and offload policies based on config. Args:

Initialize the FSDPEngine. Sets up distributed device meshes, LoRA, and offload policies based on config. Args:

Initialize the TorchTitanEngine. Sets up distributed device meshes for tensor and data parallelism, LoRA, and offload policies.

Wait for CPU worker to finish all gradient accumulation tasks.

Call parameter synchronization and asynchronous sequence generation.

broadcast tensor in tp shards across mp_group

broadcast tensor in tp shards across mp_group

Build KV cache for prefix. No grad needed.

Build kwargs dict for layer forward, based on what the layer accepts.

Replace <image> and <video> placeholder in messages with corresponding image and video which is required by processor.apply_chat_template.

Promote completed prefetches to GPU_READY.

Validate and resolve the unnormalization key for action statistics

Collect GPU buffer grads to CPU layer using K-slab flat buffer pool.

Compute FP8 alignment sizes for thd-format sequences. For FP8 block quantization, each sequence must be padded to a multiple of lcm(16, align

Compute whether data collection is needed for the current worker. This function determines whether the current worker should collect data based o

Compute the reference log prob of the batch.

Computes information about prompts and responses from a batch. This is an internal helper function that extracts masks and lengths for promp

Compute reward score for single sample.

Compute the values of the batch.

Convert regular tensor to NestedTensor, slicing according to input_ids_lengths. Args: v: Tensor to convert, shape [batch, seq_len]

Create a TIMM-based featurizer model with appropriate configurations. Args: model_id: The TIMM model ID to load

Create test DataProto for update_policy method

Recursively performs a deep comparison between two Python objects. - Handles NaN values correctly (NaN == NaN evaluates to True). - Handl

Discover model components via attribute introspection. Supports LLM and VLM models: - LLM: LLaMA, Qwen, Mistral, Phi, Gemma, GPT-2, DeepSeek,

Estimate the FLOPS of the vision encoder for Qwen3-VL

Async evict layer to CPU.

Evict a layer back to CPU.

Extract LoRA adapter tensors as a list of (name, tensor) tuples. Mimics the base_sync_done=True path — sends adapter deltas via LoadLoRAAdapt

Try multiple attribute paths, return first match or raise.

Run forward pass through all layers and return final hidden states. This is a shared helper used by both inference and training paths.

Returns: dict[str, torch.Tensor]: log_probs: (bs, response_len) if calculate_entropy is True:

Deterministically generate weights on the best available device from specs. Args: weight_specs: list of (name, shape, dtype) tuples

Get the data parallel mesh, handling hybrid/fully/replicate shard modes.

Get master addr and port for this worker group

Return the directory path for a given gen_step (one dir per step, no .pkl).

Extract a field from qat_config, supporting both dict and object-style access.

Context manager for safe session access with proper connection pooling. Yields: aiohttp.ClientSession: Session instance for makin

Initialize QAT configuration from actor.qat.

Insert a block of data from source into one task pool.

Introspect a layer's forward signature to determine accepted kwargs. Returns a set of accepted parameter names.

Check if input_amax has been initialized.