VeOmni Documentation

The Preprocessor Registry enables a decoupled data pipeline where configuration drives data transformation. The flow follows these steps:

1. Configuration: A YAML config specifies a source_name (e.g., sharegpt4v_pretrain).
2. Training Script: The script calls a transformation function (like process_sample()) which invokes conv_preprocess() using the source_name from the config.
3. Registry Lookup: The registry looks up the function decorated with the matching name.
4. Transformation: The registered preprocessor (defined in preprocess.py) transforms the raw data source into VeOmni's standardized conversation format.
5. Output: The standardized format is returned, ready for image processing and tokenization.

Config (qwen2_vl.yaml)
  └─> source_name: sharegpt4v_pretrain
       └─> Training Script (train_vlm.py)
            └─> process_sample() calls conv_preprocess("sharegpt4v_pretrain", ...)
                 └─> Registry looks up sharegpt4v_pretrain_preprocess()
                      └─> Preprocessor (preprocess.py) transforms raw data
                           └─> Returns standardized conversation format

Training Mixture-of-Experts (MoE) models with LoRA requires specific settings:

• Expert Parallelism (EP): To enable EP, set train.accelerator.ep_size > 1 and set model.ops_implementation.moe_implementation to fused_triton.
• Shared LoRA: Set model.lora_config.share_expert_lora: true for Mode 2 (one LoRA per layer).
• Auto-mapping: For MoE, gate_proj, up_proj, and down_proj automatically map to the fused expert parameters.

model:
  model_path: Qwen3-30B-A3B-merge
  ops_implementation:
    attn_implementation: flash_attention_2
    moe_implementation: eager           # EP (ep_size > 1) REQUIRES fused_triton
  lora_config:
    rank: 16
    alpha: 32
    lora_modules: [q_proj, k_proj, v_proj, o_proj, gate_proj, up_proj, down_proj]
    share_expert_lora: true                            # one LoRA per layer (Mode 2)

train:
  init_device: meta
  accelerator:
    ulysses_size: 1
    ep_size: 1                          # set >1 to enable EP; also set moe_implementation: fused_triton
    fsdp_config:
      fsdp_mode: fsdp2

The kernel selection process follows a specific lifecycle from import to runtime:

1. Import Time: apply_ops_patch() is called, which registers facade names (like Flash/Flex attention) with Sequence Parallel (SP).
2. Config Parse Time: OpsImplementationConfig.__post_init__() validates backends, rewrites attention implementations for SP compatibility, and populates the global ops config singleton via set_ops_config(self).
3. Model Build Time: build_foundation_model triggers apply_ops_config(ops). This installs the LOSS_MAPPING (binding the specific CE kernel to the loss function via functools.partial) and applies global patches (e.g., MoE, load-balancing loss). OpSlot.bind(impl_name) is used for per-model dispatch.
4. Runtime: Operations dispatch to the pre-selected kernels. For example, attention uses ALL_ATTENTION_FUNCTIONS[config._attn_implementation], and loss uses the pre-bound self.loss_function to avoid per-forward lookup overhead.

Muon's natural unit is a full matrix, but attention computes scores per head. Setting muon_head_group_size (the number of heads per block) allows splitting a head-stacked projection into row blocks, giving each block its own polar factor. This is known as "Muon Split".

When to use it

• For Wide/Square matrices (e.g., classic MHA q_proj): Splitting might actually decrease performance as it drops cross-head orthogonality.
• For Tall/Skinny matrices (e.g., MLA or low-rank up-projections like DeepSeek V4's q_b_proj): Splitting is highly beneficial. It prevents heads from sharing a single small update budget and restores a full-scale whitened update per head.

Configuration Example

To enable head-splitting, you must provide both muon_head_group_size and the specific module names in muon_head_split_modules via your config:

train:
  optimizer:
    type: muon
    muon_head_group_size: 1            # heads per block
    muon_head_split_modules: [q_b_proj]  # which projections to split

When Sequence Parallelism (SP) is enabled, the SequenceParallelCollator (appended to MainCollator) manages the lifecycle of tensors to ensure they are correctly sharded for each SP rank.

Data Transformation Steps

1. Label shifting: Shifts labels left by 1 token for next-token prediction.
2. SP padding and slicing:
   • Pad: Sequences are padded to be divisible by sp_size using sp_pad_value.
   • Slice: Sequences are sliced for the current rank using tensor.narrow(dim, rank * chunk, chunk).
3. Flash attention kwargs: Computed from position_ids before slicing to ensure the full sequence boundaries are captured.

Tensor Shapes and Properties after Collator

Default DataCollateInfo Mapping

The BaseTrainer class is the central orchestrator in VeOmni. It manages the entire training lifecycle by integrating several core components:

• Distributed Setup: Handles initialization of process groups and parallel states (e.g., DP, TP, EP).
• Component Construction: Automatically builds the model, optimizer, scheduler, and dataloaders based on provided configurations.
• Training Loop: Manages the execution of epochs and steps, including gradient accumulation.
• State Management: Handles checkpointing and resuming training.
• Extensibility: Uses a callback system to allow users to inject custom logic without modifying the core loop.

Core Attributes:

• args: Global configuration for model, data, and training.
• model: The parallelized model (wrapped with FSDP/DDP).
• optimizer & lr_scheduler: The optimization components.
• train_dataloader: The distributed data loader.
• callbacks: The handler for registered callback objects.

When using the qwen_omni_offline_av preprocessor, you provide video and audio as a single paired unit to avoid re-decoding containers during every training epoch. Each sample is a dictionary where the videos list contains paired A/V dictionaries.

Paired A/V Structure:

• frames: A List[bytes] of PNG or JPEG-encoded frames.
• audio: WAV-encoded bytes or a 1-D np.ndarray of mono samples.
• video_fps: The frame rate of the video (falls back to mm_configs.fps).
• audio_fps: The audio sample rate. If using a 1-D np.ndarray, this field is required.

Tokenization Behavior: Using a single <video> marker in the conversation binds the entire paired dictionary. The processor detects the presence of audio and automatically interleaves tokens using the standard omni path (<vision_bos><audio_bos> ... <|video_pad|> / <|audio_pad|> ... <audio_eos><vision_eos>).

{
  "videos": [
    {
      "frames": ["<png-bytes-frame-0>", "<png-bytes-frame-1>", "..."],
      "audio":  "<wav-bytes>",
      "video_fps": 2.0,
      "audio_fps": 16000
    }
  ],
  "conversations": [
    {"from": "human", "value": "<video>\nWhat is happening in the clip?"},
    {"from": "gpt",   "value": "Someone is speaking near a car."}
  ]
}

The Custom Preprocessor Registry is an extensible system for registering functions that convert raw data from specific sources into VeOmni's standardized multimodal conversation format.

Key Concepts:

• Dataset: A class responsible for loading data (e.g., MappingDataset).
• Preprocessor: A function responsible for format conversion. The registry manages these functions, not the dataset classes.
• Registration: Preprocessors are registered using the @PREPROCESSOR_REGISTRY.register("name") decorator. Once the veomni.data.multimodal module is imported, all registered preprocessors are automatically available.

Workflow:

1. Define a function with the @PREPROCESSOR_REGISTRY.register decorator.
2. Ensure the module is imported (ideally by adding it to veomni/data/multimodal/preprocess.py).
3. Access the preprocessor via the conv_preprocess convenience function or through configuration files.

@PREPROCESSOR_REGISTRY.register("my_custom_source")
def my_custom_source_preprocessor(conversations, **kwargs):
    # logic to convert raw data to VeOmni format
    return constructed_conversation