[lora, ops, dist] feat: add SwiGLU-OAI expert LoRA for M3.1-Flash

.agents/knowledge/constraints.md

@@ -119,6 +119,11 @@ Core entry points:
 8c. **Differentiable SP collectives must remain in every rank's autograd graph**
    - Rank-local packed boundaries can make a gathered halo unused on some ranks. Preserve a zero-valued dependency on the gathered tensor so every rank executes the same backward collective sequence; otherwise a later all-reduce can be matched against another rank's gather backward and fail with inconsistent collective counts.
 
+8d. **A shared MoE load-balancing loss is valid only when it matches the model's router semantics**
+   - Do not attach the Switch-style `load_balancing_loss_func` merely because a model exposes raw router logits. That loss applies `softmax(router_logits)` and performs its own top-k selection.
+   - Routers such as MiniMax M3 instead select experts from `sigmoid(router_logits) + e_score_correction_bias`, then derive expert weights from the unbiased sigmoid scores. A softmax-based auxiliary loss can therefore optimize a different expert assignment from the one executed by the model.
+   - For these models, leave router-logit capture and the generic auxiliary loss disabled until a model-specific balancing algorithm, statistics scope, and distributed update rule are verified against the model's training recipe.
+
 ## Data Pipeline
 
 Core files:
@@ -167,8 +172,9 @@ Core files:
 ### Multimodal Data
 
 15. **Multimodal preprocessing pipeline (`veomni/data/multimodal/` + `veomni/data/data_transform.py`)**
-    - The two orchestrators differ in where tokenization and label masking happen — do not assume a single shared order.
+    - The orchestrators differ in where tokenization and label masking happen — do not assume a single shared order.
       - `process_sample_qwen_vl()`: `conv_preprocess()` → `fetch_images` / `fetch_videos_metadata` → `processor.image_processor` / `processor.video_processor` for pixel features only → `chat_template.encode_messages()`, which does **both** tokenization and label masking.
+      - `process_sample_minimax_m3_vl()`: shares the image/video loading of `process_sample_qwen_vl()` (`_prepare_vlm_media_inputs`, which also calls `processor.video_processor` with metadata), then `MiniMaxM3VLChatTemplate` renders the checkpoint's native template through `processor.apply_chat_template`, tokenizes the text and masks labels turn by turn.
       - `process_sample_qwen_omni()`: takes no chat template. `conv_preprocess()` → `fetch_images/videos/audios` → `processor(text=..., images=..., videos=..., audios=...)` for tokenization → labels masked inline by a user/assistant token scan.
     - Images: load → RGB PIL → `smart_resize` (pixel min/max, scale_factor for grid alignment, max aspect ratio).
     - Videos: `torchcodec` decode → `calculate_frame_indices` (FPS, min/max frames, `frame_factor`/`frame_factor_remainder` for VAE-friendly counts); optional paired audio.

.agents/knowledge/multimodal_metadata.md

@@ -217,6 +217,7 @@ guarantees:
 | qwen2_vl | ✅ wired | Own `collate_multimodal_metadata` (non-window ViT, same formula as qwen3_vl). |
 | qwen2_5_vl | ✅ wired | Window-attention ViT. `collate_multimodal_metadata` ports `get_window_index` host-side; `get_metadata_collate_func` `partial`-closes the vision-config dims. |
 | qwen2_5_omni | ✅ wired | Same window-attention ViT as qwen2_5_vl. Also exposes `get_extra_collate_infos` (audio); `get_metadata_collate_func` is patched on the thinker, the top-level model delegates. |
+| minimax_m3_vl | ✅ wired | Own `collate_multimodal_metadata`: `{image,video}_grid_thw_list` plus per-modality `vit_{image,video}_cu_seqlens` / `max_seqlen` that keep visual sample boundaries and isolate the SP padding tail. |
 
 ## Adding the hook to a new model (checklist)
 

README.md

@@ -79,6 +79,7 @@ Our guiding principles when building VeOmni are:
 | [GPT-OSS](https://huggingface.co/openai/gpt-oss-120b)    | 120B                          | [gpt_oss_120b_lora_ep4.yaml](configs/text/gpt_oss_120b_lora_ep4.yaml) |
 | [Qwen2-3 Omni](https://huggingface.co/Qwen)              | 7B/30BA3B                     | [qwen25_omni.yaml](configs/multimodal/qwen25_omni/qwen25_omni.yaml)   |
 | [Wan](https://huggingface.co/Wan-AI)                     | Wan2.1-I2V-14B-480P           | [wan_sft.yaml](configs/dit/wan_sft.yaml)                              |
+| [MiniMax M3 VL](docs/examples/minimax_m3_vl.md)          | Checkpoint-dependent          | [minimax_m3_vl.yaml](configs/multimodal/minimax_m3_vl/minimax_m3_vl.yaml); [NPU](configs/multimodal/minimax_m3_vl/minimax_m3_vl_npu.yaml) |
 | [MiniMax H3](docs/examples/minimax_h3.md)                | Checkpoint-dependent          | [FL2VA](configs/dit/minimax_h3_fl2va_offline.yaml); [Ref2VA](configs/dit/minimax_h3_ref2va_cfg_offline.yaml) |
 | [LTX-2.3](https://huggingface.co/Lightricks/LTX-2.3)     | Checkpoint-dependent          | [ltx2_av_lora.yaml](configs/dit/ltx2_av_lora.yaml); [training guide](docs/examples/ltx-2.3.md) |
 

configs/multimodal/minimax_m3_vl/minimax_m3_vl.yaml

@@ -0,0 +1,46 @@
+model:
+  model_path: MiniMaxAI/MiniMax-M3
+  ops_implementation:
+    attn_implementation: flash_attention_2
+    moe_implementation: fused_triton
+    cross_entropy_loss_implementation: eager
+  optimizer:
+    lr: 1.0e-5
+    lr_decay_style: cosine
+  accelerator:
+    gradient_checkpointing:
+      enable_reentrant: false
+    init_device: meta
+    ulysses_size: 2
+    ep_size: 2
+    fsdp_config:
+      fsdp_mode: fsdp2
+  chat_template: minimax_m3_vl
+
+data:
+  train_path: configs/multimodal/data/tulu_sharegpt4v_llavavideo.yaml
+  data_type: conversation
+  max_seq_len: 4096
+  train_size: 80000000
+  mm_configs:
+    image_max_pixels: 602112
+    video_max_pixels: 602112
+    video_total_pixels: 2889523
+    max_frames: 4
+    fps: 2.0
+    use_audio_in_video: false
+
+train:
+  wandb:
+    enable: false
+    project: VeOmni
+    name: MiniMax-M3-vl-sft
+  freeze_vit: false
+  num_train_epochs: 1
+  micro_batch_size: 1
+  global_batch_size: 4
+  max_steps: 500
+  checkpoint:
+    output_dir: MiniMax-M3-vl-sft
+    manager: dcp
+    save_hf_weights: false

configs/multimodal/minimax_m3_vl/minimax_m3_vl_npu.yaml

@@ -0,0 +1,47 @@
+model:
+  model_path: MiniMaxAI/MiniMax-M3
+  ops_implementation:
+    attn_implementation: flash_attention_2
+    moe_implementation: fused_npu
+    rms_norm_implementation: torch_npu
+    cross_entropy_loss_implementation: eager
+  optimizer:
+    lr: 1.0e-5
+    lr_decay_style: cosine
+  accelerator:
+    gradient_checkpointing:
+      enable_reentrant: false
+    init_device: meta
+    ulysses_size: 2
+    ep_size: 2
+    fsdp_config:
+      fsdp_mode: fsdp2
+  chat_template: minimax_m3_vl
+
+data:
+  train_path: configs/multimodal/data/tulu_sharegpt4v_llavavideo.yaml
+  data_type: conversation
+  max_seq_len: 4096
+  train_size: 80000000
+  mm_configs:
+    image_max_pixels: 602112
+    video_max_pixels: 602112
+    video_total_pixels: 2889523
+    max_frames: 4
+    fps: 2.0
+    use_audio_in_video: false
+
+train:
+  wandb:
+    enable: false
+    project: VeOmni
+    name: MiniMax-M3-vl-sft
+  freeze_vit: false
+  num_train_epochs: 1
+  micro_batch_size: 1
+  global_batch_size: 4
+  max_steps: 500
+  checkpoint:
+    output_dir: MiniMax-M3-vl-sft
+    manager: dcp
+    save_hf_weights: false

docs/design/op_selection.md

@@ -22,6 +22,7 @@ If a field is not in this table, it is not an op-selection knob.
 | DSA attention | `dsa_attention_implementation` | `eager`, `flashmla_cudnn` (GLM-DSA), `tilelang` (DeepSeek-V4), `triton` (DeepSeek-V4), `cann` (DeepSeek-V4) | `"eager"` | Model `__init__` via an instance-local `VeomniOp` |
 | mHC | `mhc_implementation` | `eager`, `tilelang` (DeepSeek-V4, SM90+), `cann` (DeepSeek-V4, Ascend) | `"eager"` | Model `__init__` via three instance-local `VeomniOp`s (`pre`, `post`, `head`) |
 | QSA attention | `qsa_attention_implementation` | `eager`, `tilelang` (Qwen4-Exp, NVIDIA CUDA), `torch_npu` (Qwen4-Exp, Ascend NPU) | `"eager"` | Model `__init__` via an instance-local `VeomniOp` |
+| MiniMax sparse attention | — (no config field until a kernel backend exists) | `eager` | `"eager"` | Model `__init__` via two instance-local `VeomniOp`s (`indexer`, `attention`) of op `minimax_sparse_attention`, on sparse layers only (dense layers bind the standard `attention` op from `attn_implementation`); the eager attention runs `torch_npu.npu_fusion_attention` on NPU tensors |
 | Cross-entropy loss | `cross_entropy_loss_implementation` | `eager`, `liger_kernel`, `chunk_loss` | `"liger_kernel"` (GPU, MLU), `"chunk_loss"` (NPU) | Model `__init__` via an instance-local `VeomniOp` |
 | RMSNorm | `rms_norm_implementation` | `eager`, `liger_kernel`, `torch_npu`, `triton` (per-model; DeepSeek-V3) | `"liger_kernel"` (GPU, MLU), `"torch_npu"` (NPU) | Model `__init__` via an instance-local `VeomniOp` |
 | SwiGLU MLP | `swiglu_mlp_implementation` | `eager`, `liger_kernel` | `"liger_kernel"` (GPU, MLU), `"eager"` (NPU) | Model `__init__` via an instance-local `VeomniOp` |
@@ -518,6 +519,12 @@ separate `ModuleList` fork. LoRA experts follow the wrapped module's
 `veomni_moe.impl` and fall back to `eager` when that `moe_experts_lora` row is
 unavailable on the device.
 
+MiniMax M3 experts use the `swiglu_oai` variant, which registers `eager`,
+`fused_triton` and `fused_npu` only. It shares the standard merged-fc1 Triton
+and NPU paths (including EP) with a `swiglu_alpha` argument that replaces SiLU
+with the clamped `(up + 1) * gate * sigmoid(alpha * gate)` activation; selecting
+an implementation without a row, such as `fused_quack`, fails at model build.
+
 DeepSeek-V4 keeps eager DSA indexer and attention as its defaults, with optional
 SM90+ `tilelang` and Ascend `cann` indexer and attention implementations, plus
 ROCm `triton` attention. Its MoE path uses the independent
@@ -780,3 +787,38 @@ Temporary packing storage trades memory traffic for fewer collective launches.
 `tests/models/test_qwen4_exp_gdn_exchange.py` covers exact CPU/Gloo permutations,
 projection view storage, and native NCCL/HCCL model parity with same/mixed dtypes.
 CPU tests do not establish native throughput or numerical acceptance.
+
+
+## SwiGLU OAI expert LoRA selection
+
+Native MoE LoRA recognizes a base operator with `variant="swiglu_oai"` and
+selects `shared_swiglu_oai` or `independent_swiglu_oai` in `moe_experts_lora`.
+Both expose eager and NPU implementations. The model supplies positive finite
+`swiglu_alpha` and `swiglu_limit`; unsupported activations still fail before
+the wrapper takes ownership of parameters. No generated model file is edited.
+
+The NPU variants reuse the standard expert dispatch, grouped GEMMs and combine
+operations, with the configured clamped activation after gate/up adapter deltas.
+CPU-emulated composition tests cover the tensor math and gradients; real NPU
+and EP execution remain separate validation requirements.
+
+`VeOmniLoraModel.disable_adapter()` also covers native experts that expose a
+base `VeomniOp` handle. Expert forward then calls that handle with the original
+base weights and activation attributes and does not read adapter tensors.
+The context preserves prior disable states across nesting and exceptions.
+An expert without a base operator fails before any adapter state changes.
+
+
+The base `moe_experts/swiglu_oai` row serves MiniMax M3.1-Flash experts.
+Clamped standard experts are explicitly rejected until their adapter activation
+has a matching implementation, preventing an incorrect unclamped reference.
+
+For large frozen-base adapters, `build_parallelize_model(..., is_peft_model=True,
+preserve_frozen_dtype=True)` keeps caller-selected base and buffer dtypes while
+upcasting trainable floating parameters to FP32. This opt-in does not choose
+which base weights need FP32; the caller must prepare that policy before sharding.
+This policy requires FSDP2 with enabled mixed precision and a non-None
+`param_dtype`: its parameter casting supplies
+matching compute dtypes. DDP and unsharded execution are rejected because the
+frozen BF16 base and FP32 adapters otherwise have incompatible input dtypes.
+The default remains whole-model FP32 conversion under mixed precision.