VeOmni Test Suite Overview

VeOmni Test Suite Overview#

This document surveys all tests in the VeOmni project, describes their purpose and organization, and provides guidance on which tests to add when onboarding a new model.

Directory Structure#

tests/
├── tools/                          # Shared test infrastructure (comparison, data gen, launch)
├── toy_config/                     # Minimal model configs for fast CI testing
├── testdata/                       # Sample images, audio, etc.
│
├── models/                         # Single-GPU model correctness
│   ├── test_models_patch.py        # Fwd/bwd across attention & MoE backends
│   ├── test_vlm_trainer.py         # VLM freeze_vit smoke test
│   ├── test_model_registry.py      # Model loader registry (HF vs VeOmni)
│   ├── test_checkpoint_tensor_converter.py  # Checkpoint tensor conversion (e.g. Qwen3MoE fuse)
│   ├── test_deepseek_v4_fused_moe.py  # DeepSeek-V4 fused MoE swiglu_limit plumbing
│   ├── test_padded_packed_loss.py   # Padded vs packed (cu_seqlens) loss equivalence
│   ├── test_models_logits_equal_v5.py  # HF↔VeOmni logits through the real loader
│   └── utils.py                    # ModelMode, prepare_model_modes, prepare_data
│
├── ops/                            # Fused kernel correctness & performance
│   ├── test_fused_moe_split_vs_merged.py   # Split vs merged MoE fc1
│   ├── test_quack_fused_moe.py             # Quack GEMM MoE (SM90+)
│   ├── test_fused_load_balancing_loss.py    # Triton load-balancing loss kernel
│   ├── test_deepseek_v4_kernels.py           # DeepSeek-V4 TileLang guards and numerical parity
│   ├── test_mhc_tile_kernels.py              # DeepSeek-V4 TileKernels mHC dispatch and parity
│   ├── test_flash_attn_varlen_padding.py   # Flash-attn variable-length padding
│   ├── test_seqcls_loss.py                 # Sequence classification loss
│   └── test_comp.py                        # Position embedding computation
│
├── data/                           # Data loading & preprocessing
│   ├── test_datasets.py            # Dataset loading, filtering, schema validation
│   ├── test_collators.py           # MainCollator, cu_seq_lens generation
│   ├── test_dataloader.py          # DataLoader batching
│   ├── test_dpo_data_processor.py  # DPO data processing
│   ├── test_dynamic_batching_dataset.py  # Dynamic batching by seq length
│   ├── test_prepare_fa_kwargs.py   # Flash-attn parameter construction
│   ├── test_preprocessor.py        # Token mapping, special tokens
│   ├── test_classification_data_processor.py  # Classification data processing
│   └── multimodal/
│       ├── test_vlm_data_process.py   # VLM data pipeline (HF processor vs VeOmni)
│       └── test_video_utils.py        # Video/audio loading & frame extraction
│
├── parallel/                       # Parallelism primitives
│   ├── ulysses/                    # Sequence parallelism (Ulysses)
│   │   ├── test_ulysses.py             # Basic SP attention (4+ GPUs)
│   │   ├── test_async_ulysses.py       # Async overlapping comm (4+ GPUs)
│   │   ├── test_async_ulysses_dit.py   # DiT + async SP
│   │   ├── test_qwen3_5_gated_deltanet_ulysses.py  # Gated DeltaNet + SP
│   │   ├── test_slice_input_tensor.py  # Input slicing utilities
│   │   ├── test_all_gather.py          # All-gather collective ops
│   │   └── utils.py                    # SequenceParallelTest base class
│   └── encoder_data_balance/
│       ├── test_balance_reverse.py        # Balance/recovery precision (8 GPUs)
│       └── test_balance_sorting_algo.py   # Post-MBS data sorting (CPU)
│
├── distributed/                    # Multi-GPU training correctness
│   ├── test_fsdp_equivalence.py         # Single-GPU vs FSDP2 grad equivalence
│   └── test_dummy_forward.py            # Asymmetric multimodal forward (NCCL hang prevention)
│
├── e2e/                            # End-to-end training integration
│   ├── test_e2e_parallel.py             # SP/EP parallel alignment across models
│   ├── test_e2e_training.py             # Real-model SFT smoke test (8 GPUs)
│   ├── test_e2e_training_no_reshard.py  # FSDP2 no-reshard mode
│   ├── exec_scripts.py                  # Shell command generators for real models
│   └── utils.py                         # prepare_exec_cmd, parse_training_log, ParallelMode
│
├── train_scripts/                  # Standalone trainer scripts (invoked via torchrun, not pytest)
│   ├── train_text_test.py               # Test trainer for text models
│   ├── train_vlm_test.py                # Test trainer for VLM models
│   └── train_dit_test.py                # Test trainer for DiT models
│
├── checkpoints/                    # Checkpoint save/load
│   ├── test_checkpoint_callback.py          # Callback _last_saved_step correctness
│   ├── test_trainer_saveload.py             # DCP + HF checkpoint save/load (8 GPUs)
│   ├── checkpoint_verification_utils.py     # DCP-to-HF conversion verification
│   └── utils.py                             # Command/config builders for ckpt tests
│
├── utils/                          # Misc utility tests
│   ├── test_count_flops.py                       # FLOPs estimation
│   ├── test_extra_parallel_clip_grad_norm.py      # Grad clipping with EP/EMB dims (8 GPUs)
│   ├── test_helper.py                             # EnvironMeter utility (8 GPUs)
│   ├── test_model_loader.py                       # Model loading (4 GPUs)
│   ├── test_npu_setup.py                          # NPU environment validation
│   ├── test_rank0_load_and_broadcast_weights.py   # Rank-0 load & broadcast (2+ GPUs)
│   └── test_save_safetensor_utils.py              # Safetensor save utilities (CPU)
│
└── special_sanity/
    └── check_device_api_usage.py    # CI lint: no direct .cuda / "cuda" calls

Test Categories at a Glance#

Category

Directory

GPU Req

Execution

Purpose

Model patch

tests/models/

1 GPU

pytest

Fwd/bwd correctness across attn/MoE backends

Ops / kernels

tests/ops/

0-1 GPU (SM90+ for Quack, DeepSeek-V4 TileLang, and mHC TileKernels)

pytest

Fused kernel guards, dispatch, correctness, and performance

Data pipeline

tests/data/

0-1 GPU

pytest

Data loading, collation, preprocessing

Parallelism

tests/parallel/

4-8 GPUs

torchrun / pytest

SP, EP, data-balance primitives

FSDP correctness

tests/distributed/

2+ GPUs

torchrun (subprocess + mp.spawn)

Single-GPU vs FSDP2 equivalence, dummy forward

E2E parallel

tests/e2e/

4+ GPUs

torchrun (subprocess)

SP/EP alignment across full training runs

Checkpoints

tests/checkpoints/

0-8 GPUs

pytest + torchrun

Save/load, DCP→HF conversion

Utilities

tests/utils/

0-8 GPUs

pytest + torchrun

FLOPs, grad clipping, weight broadcast

Sanity

tests/special_sanity/

0

script

Device API lint

Shared Test Infrastructure (tests/tools/)#

All shared, cross-cutting utilities live in tests/tools/:

Module

Exports

Description

comparison_utils

TensorComparator, assert_close, assert_exact, compare_metrics, print_comparison_table

Numerical comparison with tolerances; rich table output

data_generators

DummyDataset

Generates parquet dummy datasets for all modalities (text, VLM, omni, DiT)

launch_utils

find_free_port, torchrun

Port discovery; mp.spawn-based distributed launcher

training_utils

ParallelConfig, build_torchrun_cmd, materialize_weights, run_training_config, release_device_memory

Torchrun command builder, model weight materialization, training runner

Additional per-directory helpers:

File

Scope

Key Exports

tests/models/utils.py

Model patch tests

ModelMode, prepare_model_modes, prepare_data

tests/models/test_checkpoint_tensor_converter.py

Model loading

Runtime checkpoint layout conversion and fused-expert weight mapping

tests/e2e/utils.py

E2E tests

prepare_exec_cmd, parse_training_log, ParallelMode

tests/checkpoints/utils.py

Checkpoint tests

Command/config builders for trainer save/load

tests/parallel/ulysses/utils.py

SP tests

SequenceParallelTest base class, sync_tensor

Detailed Test Descriptions#

1. Model Patch Tests (tests/models/test_models_patch.py)#

Purpose: Verify that VeOmni’s patched models produce identical loss and grad_norm to HuggingFace reference across all backend combinations.

What it compares (cartesian product):

Dimension

Values

Modeling backend

HuggingFace, VeOmni

Attention implementation

eager, flash_attention_2, flash_attention_3, veomni_flash_attention_2_with_sp, veomni_flash_attention_3_with_sp

MoE implementation

eager, hardware fused backend (fused_triton on GPU, fused_npu on NPU where supported)

Liger kernel

True, False (VeOmni only)

Models covered:

  • Text / MoE: llama3_1, qwen2, qwen3_5, qwen3_5_moe, seed_oss, deepseek_v3, deepseek_v4

  • VLM: qwen2_vl, qwen2_5_vl, qwen3_vl, qwen3_vl_moe

  • Omni: qwen2_5_omni, qwen3_omni_moe

GPU: 1 GPU, runs serially per model mode.

DeepSeek-V4’s fused-MoE-specific merged gate_up_proj and swiglu_limit forwarding are covered by tests/models/test_deepseek_v4_fused_moe.py (CPU). Its kernel package import behavior, hardware guards, BF16/FP32 utility, and TileLang DSA indexer/attention numerical checks are covered by tests/ops/test_deepseek_v4_kernels.py. The guard and utility cases run on CPU; optimized numerical tests require TileLang on an SM90+ NVIDIA GPU. Registry binding plus mHC pre/post/head forward and backward parity are covered by tests/ops/test_mhc_tile_kernels.py, which requires TileKernels on an SM90+ NVIDIA GPU for kernel execution.

2. VLM Trainer Test (tests/models/test_vlm_trainer.py)#

Purpose: Smoke test that freeze_vit=True/False correctly freezes/unfreezes the vision tower.

Models: qwen2_vl, qwen3_5, qwen3_5_moe, qwen2_5_vl, qwen3_vl, qwen3_vl_moe

GPU: CPU only (builds model but no forward pass).

3. Model Registry Test (tests/models/test_model_registry.py)#

Purpose: Verify that get_model_config/class/processor returns the correct HF or VeOmni module.

GPU: CPU only.

4. Checkpoint Tensor Converter (tests/models/test_checkpoint_tensor_converter.py)#

Purpose: Test checkpoint tensor conversion protocol (e.g., Qwen3MoE expert weight fusion: per-expert → stacked gate_up_proj).

GPU: CPU only.

5. Padded vs Packed Loss (tests/models/test_padded_packed_loss.py)#

Purpose: Verify that padded input and packed input (with cu_seqlens) produce identical loss.

GPU: 1 GPU (requires flash-attn).

6. FSDP Equivalence (tests/distributed/test_fsdp_equivalence.py)#

Purpose: Verify that FSDP2 sharding produces the same grad_norm as single-GPU training (no parallelism). This catches FSDP wrapping bugs that silently corrupt gradients.

How it works:

  1. Materialize random weights from toy config

  2. Run single-GPU training (nproc=1, no FSDP)

  3. Run FSDP2 training (nproc=2+, init_device=meta)

  4. Assert grad_norm matches (loss may differ due to micro-batch splitting)

Models: qwen3, qwen3_moe, llama3.1, qwen3_5, qwen3_5_moe

GPU: 2+ GPUs.

7. Dummy Forward (tests/distributed/test_dummy_forward.py)#

Purpose: Verify that asymmetric multimodal batches (some ranks text-only, others with images/video/audio) don’t cause NCCL hangs under FSDP2. Tests that dummy_forward() is correctly invoked so all ranks participate in FSDP collectives.

Models:

  • VLM: qwen2_5_vl, qwen3_vl, qwen3_vl_moe

  • Omni: qwen2_5_omni, qwen3_omni_moe

GPU: 2 GPUs.

8. E2E Parallel Alignment (tests/e2e/test_e2e_parallel.py)#

Purpose: Full torchrun training runs across SP/EP configurations. Asserts that loss and grad_norm match regardless of parallelism settings.

Configurations tested:

  • sp_size in [1, 2], ep_size in [1] (base) or [1, 2] (MoE)

  • FSDP2 always enabled, nproc_per_node = sp_size * 4

  • 2 epochs, 2 max_steps per run

Models: All supported text, VLM, omni, and DiT models.

GPU: 4+ GPUs (up to 8 for sp=2).

9. E2E Training Smoke Tests (tests/e2e/test_e2e_training*.py)#

Purpose: Smoke tests with real model weights (qwen3_0p6b_base + Tulu-3 SFT dataset). Validates that training completes without errors.

  • test_e2e_training.py — standard FSDP2 training (8 GPUs)

  • test_e2e_training_no_reshard.py — FSDP2 no-reshard mode (8 GPUs)

10. Checkpoint Save/Load (tests/checkpoints/)#

Test

Purpose

GPU

test_checkpoint_callback.py

_last_saved_step state tracking

CPU

test_trainer_saveload.py

DCP + HF checkpoint formats, resume training

8 GPUs

11. Ops / Kernel Tests (tests/ops/)#

Test

Purpose

GPU

test_fused_moe_split_vs_merged.py

Split vs merged fc1 in fused MoE

1 GPU

test_quack_fused_moe.py

Quack GEMM MoE backend

SM90+

test_kernel_registry_numerical.py

Numerical alignment per (op, variant, impl)

CUDA; the FlashQLA chunk_gated_delta_rule case skips unless running on SM90 (Hopper) — SM10x WIP upstream.

test_fused_load_balancing_loss.py

Triton load-balancing loss

CUDA

test_deepseek_v4_kernels.py

CPU import/hardware guards plus TileLang DSA numerical parity

CPU for guards; TileLang + NVIDIA SM90+ for optimized kernels

test_mhc_tile_kernels.py

TileKernels mHC registry dispatch and forward/backward parity

TileKernels + NVIDIA SM90+

test_flash_attn_varlen_padding.py

Flash-attn variable-length padding

CUDA

test_seqcls_loss.py

Sequence classification loss

CUDA (optional)

test_comp.py

Position embedding computation

CUDA

12. Parallelism Primitive Tests (tests/parallel/)#

Test

Purpose

GPU

test_ulysses.py

Basic Ulysses SP attention

4+

test_async_ulysses.py

Async overlapping communication

4+

test_async_ulysses_dit.py

DiT + async SP

4+

test_qwen3_5_gated_deltanet_ulysses.py

Gated DeltaNet + SP

4+

test_slice_input_tensor.py

SP input slicing utilities

CPU

test_all_gather.py

All-gather collective ops

multi

test_balance_reverse.py

Encoder data balance recovery

8

test_balance_sorting_algo.py

Post-MBS sorting algorithm

CPU

New Model Onboarding: Test Checklist#

When adding a new model to VeOmni, the following tests should be created or updated. See also: Testing a New Model for Transformers v5 for step-by-step instructions.

Required Tests#

Step

Test File

What to Do

1. Create toy config

tests/toy_config/<model>_toy/

Minimal config (few layers, small dims). Add README.md noting the source config and changes.

2. Model patch (fwd/bwd)

tests/models/test_models_patch.py

Add a pytest.param(...) entry to the model parametrize. Filter unsupported attn/MoE modes if needed.

3. E2E parallel alignment

tests/e2e/test_e2e_parallel.py

Add entry to text_test_cases (text) or the appropriate VLM/omni list. Set max_sp_size=1 if SP not yet supported.

4. FSDP equivalence

tests/distributed/test_fsdp_equivalence.py

Add entry to verify single-GPU vs FSDP2 grad_norm matches.

Conditional Tests (depending on model type)#

Condition

Test File

What to Do

VLM model

tests/models/test_vlm_trainer.py

Add toy config to _FREEZE_VIT_VLM_CASES_*.

VLM model

tests/distributed/test_dummy_forward.py

Add test case for asymmetric multimodal batches.

MoE model

tests/models/test_models_patch.py

Set is_moe=True to test eager vs fused MoE backends.

MoE model

tests/e2e/test_e2e_parallel.py

Set is_moe=True to include ep_size iteration.

MoE with fused experts

tests/models/test_checkpoint_tensor_converter.py

Add converter tests if a custom CheckpointTensorConverter is needed.

Custom checkpoint layout

tests/models/test_checkpoint_tensor_converter.py

Add converter tests for any on-disk HF↔VeOmni key or tensor-layout conversion.

Custom fused kernels

tests/ops/

Add kernel-specific correctness tests.

New data modality

tests/data/

Add data processing and collation tests.

Verification Commands#

# Collect test cases for the new model
pytest --collect-only -k <model_name>

# Run single-GPU model patch test
pytest tests/models/test_models_patch.py -k <model_name>

# Run VLM freeze test (VLM only)
pytest tests/models/test_vlm_trainer.py -k <model_name>

# Run FSDP equivalence (2+ GPUs)
pytest tests/distributed/test_fsdp_equivalence.py -k <model_name>

# Run E2E parallel alignment (4+ GPUs)
pytest tests/e2e/test_e2e_parallel.py -k <model_name>

Test Execution Flow#

Model Patch Test Flow#

pytest → test_models_patch_fwd_bwd(config, is_moe, ...)
  → prepare_model_modes(is_moe) → [(HF, eager), (HF, fa2), (VeOmni, fa2_sp), ...]
  → for each mode:
      apply_ops_config(mode-specific OpsImplementationConfig)
      BaseTrainer._build_model() → build_foundation_model(config, ops_implementation=...)
      TrainerTest.forward_backward_step(dummy_batch)
      → record loss, grad_norm
  → compare_multi_items(all_results, rtol, atol)

E2E Parallel Test Flow#

pytest → test_text_parallel_align(model_name, config, ...)
  → materialize_weights(config) → random weights on disk
  → DummyDataset(dataset_type) → parquet files
  → for each (sp_size, ep_size):
      prepare_exec_cmd() → torchrun command
      subprocess.run(torchrun ... train_text_test.py ...)
        → TestTextTrainer.train() → log_dict.json
  → compare_multi_items(all_log_dicts, rtol, atol)

FSDP Equivalence Test Flow#

pytest → test_text_fsdp_equivalence(config, ...)
  → materialize_weights(config)
  → DummyDataset(text)
  → run_training_config(nproc=1, init_device=device)      # baseline
  → run_training_config(nproc=2+, init_device=meta, fsdp2) # FSDP
  → compare_metrics(baseline_grad_norm, fsdp_grad_norm)

Architecture Notes#

Resolved consolidations#

The following redundancies have been addressed:

  • Shared training utils centralized: ParallelConfig, build_torchrun_cmd, materialize_weights, run_training_config, and release_device_memory now live in tests/tools/training_utils.py. Both tests/e2e/ and tests/distributed/ import from tests/tools — no cross-directory imports between test subdirectories.

  • ModelMode naming conflict resolved: The e2e parallelism dataclass was renamed to ParallelMode (sp_size, ep_size) to distinguish it from ModelMode in tests/models/utils.py (modeling_backend, attn_implementation, …). ParallelConfig in tests/tools/training_utils.py adds fsdp_mode on top.

  • distributed_test_helpers.py removed: Shared helpers moved to tests/tools/training_utils.py; tests/distributed/ tests import directly from tests/tools.

  • Thin wrappers removed: compare_multi_items / print_all_values wrappers in tests/e2e/utils.py have been replaced with direct imports of compare_metrics / print_comparison_table from tests.tools.

  • Train scripts separated: train_text_test.py, train_vlm_test.py, and train_dit_test.py are standalone trainer scripts (not pytest tests). They have been moved from tests/e2e/ to tests/train_scripts/ to clarify their role.

Remaining items for future work#

  • tests/models/utils.py has its own compare_multi_items / print_all_values with custom table formatting based on ModelMode fields. These are not simple wrappers and serve a different purpose from tests.tools.compare_metrics.

  • tests/e2e/test_e2e_training.py uses real model weights and exec_scripts.py, while test_e2e_parallel.py uses toy configs and prepare_exec_cmd. These serve different purposes (smoke test vs equivalence) but the naming doesn’t reflect this.