[BREAKING][ops, model] feat: GPU-optimal ops defaults + strict NPU validation

[TimYangst]

What does this PR do?

Flip OpsImplementationConfig per-op defaults from "eager" to GPU-optimal
(Liger / Triton / fused_triton), and tighten the validator + per-model
dispatch so misconfigurations raise loudly instead of silently downgrading.

Default changes (GPU):

Field	Old	New
attn_implementation	flash_attention_2	unchanged
moe_implementation	eager	fused_triton
cross_entropy_loss_implementation	chunk_loss	liger_kernel
rms_norm_implementation	eager	liger_kernel
swiglu_mlp_implementation	eager	liger_kernel
rotary_pos_emb_implementation	eager	liger_kernel
load_balancing_loss_implementation	eager	triton

NPU users must now pin every per-op field explicitly — to an
NPU-supported value (npu / chunk_loss / fused_npu / triton) or to
eager when the op has no NPU backend. The validator raises with a
per-op allow-list message at config-parse time.

Builds on #678 (registration-based kernel framework) and #708
(chunk_loss reshape fix).

API and Usage Example

NPU: pin every per-op field. Minimal NPU-friendly config:

model:
  ops_implementation:
    attn_implementation: flash_attention_2
    moe_implementation: fused_npu
    cross_entropy_loss_implementation: chunk_loss
    rms_norm_implementation: npu
    rotary_pos_emb_implementation: npu
    swiglu_mlp_implementation: eager           # no NPU backend
    load_balancing_loss_implementation: eager  # triton-ascend not exposed as `triton`

GPU: existing YAMLs that don't override the per-op fields now pick up
the GPU-optimal kernels automatically. The gpu extra ships
liger-kernel + triton, so no setup change is needed.

Models with structurally-incompatible kernels (Wan rope_apply,
Qwen2-VL multimodal RoPE NPU) opt out via extra_backends[name] = None.
The validator raises "explicitly disabled for {model}" and points the
user at the allowed values. Wan's three YAMLs ship pinned to eager.

Design & Code Changes

Defaults & validation — veomni/arguments/arguments_types.py,
veomni/ops/kernels/{rms_norm,rotary,swiglu,load_balancing_loss}/__init__.py

• Flip per-op dataclass defaults; sync each OpSpec.default to match.
• Rewrite _validate_implementations:
  1. Check NPU compatibility against an explicit _NPU_ALLOWED allow-list.
  2. Gate load_balancing_loss=triton on is_package_available("triton").
  3. Hard-assert at validation time that every registered op
     (list_ops()) has an entry in _NPU_ALLOWED — catches future op
     additions that miss the map.
• Per-package availability for liger_kernel / torch_npu / triton is
  checked at the resolution sites (_check_requires in
  apply_per_model_patches / apply_global_ops), not duplicated in the
  validator.

Trade-off — hardcoded allow-list: _NPU_ALLOWED is hand-curated rather
than derived from BackendSpec.requires. The name "triton" is reused
with different hardware semantics (load-balancing loss runs on NPU via
triton-ascend, but DeepSeek-V3's batch-invariant RMSNorm /
deterministic RoPE Triton kernels are mainline-CUDA-only), so a
requires-based inference would conflate them. Third-party NPU backends
still extend cleanly via per-model extra_backends in device_patch.py;
the allow-list only gates the hardware-class names exposed on the public
dataclass.

Strict per-model dispatch — veomni/ops/config/registry.py

• apply_per_model_patches raises on every missing backend except
  eager. extra_backends[name] = None is repurposed as an
  explicit-raise opt-out for cases where the global registry default
  would silently bind a wrong-signature kernel; the available-list error
  excludes the disabled entry.

Safe-fallback path — veomni/models/auto.py

• Replace the silent apply_ops_config(OpsImplementationConfig())
  fallback (when callers omit ops_implementation) with
  OpsImplementationConfig.all_eager() — every per-op field pinned to
  "eager". attn_implementation keeps the dataclass default
  (flash_attention_2) when flash-attn is importable (the common case
  on production GPU hosts), falling back to "eager" when it isn't, so
  CPU / minimal-deps environments don't crash on import-time FA loading.
  Standalone scripts (inference, weight materialization, dummy-forward
  tests) keep working on every accelerator.

Per-model device patches

• qwen3_vl/device_patch.py — drop historical liger_kernel: None
  opt-outs (verified standard-shape RMSNorm + RoPE; Liger drops in
  cleanly).
• wan/device_patch.py — keep rotary_pos_emb: liger_kernel: None as
  the explicit-raise marker; move Wan's Triton RoPE from the
  _custom_wan callback into
  extra_backends["rotary_pos_emb"]["triton"] so it survives the new
  strict-raise gate (a callback can't run after the gate rejects unknown
  backends).
• qwen2_vl/device_patch.py — keep rotary_pos_emb: npu: None as the
  explicit-raise marker for the multimodal RoPE NPU mismatch.
• configs/dit/wan*.yaml (3 files) — pin
  rotary_pos_emb_implementation: eager since Wan's rope_apply can't
  host the new liger_kernel default.

Test infrastructure — tests/tools/training_utils.py and downstream

• New resolve_ops_overrides(model_name) helper emits hardware-aware
  --model.ops_implementation.X=Y flags; npu_skip_marker(model_name)
  skip factory for models with no NPU+eager path.
  _NPU_PER_MODEL_OVERRIDES covers DeepSeek-V3 (no NPU RMSNorm/RoPE)
  and Qwen2-VL / Qwen2.5-VL / Qwen2.5-Omni (multimodal RoPE).
• model_name threaded through build_torchrun_cmd /
  run_training_config and into tests/e2e/utils.py,
  tests/distributed/test_fsdp_equivalence.py, tests/checkpoints/utils.py,
  tests/data/{test_datasets,test_multisource_dataset,test_dynamic_batching_dataset}.py.
• tests/models/test_padded_packed_loss.py, tests/ops/test_moe_hw_gate.py
  — pin all-eager kwargs so tests don't couple to the new GPU defaults.

Docs — docs/usage/arguments.md, docs/design/kernel_selection.md

• Refresh defaults, NPU validation contract, lifecycle "Ownership"
  paragraph (now describes the all_eager() fallback), per-section YAML
  examples; replace VEOMNI_USE_LIGER_KERNEL references in the
  transformers-v5 comparison table with the corresponding
  OpsImplementationConfig field name.

Test

• make quality clean.
• pytest tests/ops/test_moe_hw_gate.py — 15/15 pass.
• GPU host: constructed OpsImplementationConfig() with new defaults —
  validates clean; logs show liger_kernel / fused_triton / triton
  bound for every op.
• Mocked is_torch_npu_available()=True: every GPU-only choice
  (Liger rms_norm/rotary/swiglu, Liger CE, Triton RMSNorm via
  DeepSeek-V3, fused_triton MoE) raises with the listed alternatives.
  All-NPU config (fused_npu + chunk_loss + npu + eager +
  triton) passes.
• Mocked GPU runtime: fused_npu / npu rms_norm raise with
  platform-mismatch errors.
• Per-model strict-raise: simulated apply_per_model_patches for
  Wan-shape extra_backends confirms rotary_pos_emb=liger_kernel
  (default) raises with "explicitly disabled for Wan" +
  available=['eager','npu','triton']; explicit rotary_pos_emb=eager
  succeeds.
• Verified Qwen3-VL Qwen3VLTextRMSNorm + apply_rotary_pos_emb are
  HF-standard on transformers v4.57.3 and v5 — the liger_kernel lift
  is safe.
• Pre-existing tests/ops/test_seqcls_loss.py failure (Triton driver
  missing on this CPU dev box) reproduces on main — unrelated.

Pending CI:

• gpu_e2e_test.yml, gpu_unit_tests.yml — should pass with the gpu
  extra (ships liger-kernel + flash-attn).
• npu_e2e_test.yml, npu_unit_tests.yml — resolve_ops_overrides
  threaded through every torchrun in tests/{e2e,distributed,checkpoints,data}
  emits NPU-supported overrides per model.

Checklist Before Submitting

• Read the Contribute Guide
• Applied pre-commit checks
• Added/updated documentation
• Added tests to CI workflow (or explained why not feasible) —
  resolve_ops_overrides covers every model NPU CI exercises;
  test_moe_hw_gate.py pins all-eager. No new standalone test file: the
  change is config-default + validator behavior, exercised end-to-end by
  every tests/e2e/test_e2e_parallel.py parametrization.

[gemini-code-assist]

Code Review

This pull request updates the default operator implementations to be GPU-optimal (switching from 'eager' to 'liger_kernel' or 'triton') and introduces a strict validation layer to ensure selected backends are compatible with the host hardware, particularly for Ascend NPU. It also implements a conservative all-eager fallback for standalone scripts to maintain portability and updates the test suite with hardware-aware override logic. Feedback was provided regarding the hardcoded NPU allow-list, which currently restricts third-party backend registration without framework modifications, contradicting stated design goals.

[TimYangst]

@gemini-code-assist please help review the pr

[gemini-code-assist]

Code Review

This pull request transitions the default kernel implementations to GPU-optimal backends like Liger and Triton, while providing an "all_eager" configuration for hardware-agnostic fallbacks. It also refactors NPU compatibility validation and updates model-specific patches for non-standard architectures. Feedback indicates that the "all_eager" method should explicitly include "attn_implementation" to avoid Flash Attention dependencies in safe environments, and that a missing validation check for NPU compatibility tables should be implemented as originally intended.