[veomni] feat: bump veomni to v0.1.8

- Fix parallel_state init parameter (ep_size -> extra_parallel_sizes) and use set type for basic_modules
- Automatically rewrites `flash_attention_2/3/4` to VeOmni SP-aware variants
- Add _prepare_veomni_flash_attention_kwargs to precompute cu_seq_lens for packed sequences
- Slice position_ids when sp_enabled

Author: deerlu

.github/workflows/e2e_ppo_trainer_veomni_vllm.yml

@@ -117,7 +117,8 @@ jobs:
         run: |
           pip3 install -r requirements-test.txt
           pip3 install --no-deps -e .
-          pip3 install git+https://github.com/ByteDance-Seed/VeOmni.git@v0.1.4
+          pip3 install git+https://github.com/ByteDance-Seed/VeOmni.git@v0.1.8 --ignore-requires-python --no-deps
+          pip3 install transformers==4.57.3
       - name: Prepare GSM8K dataset
         run: |
           ray stop --force

.github/workflows/e2e_sft_llm.yml

@@ -105,7 +105,8 @@ jobs:
           pip3 install peft
           pip3 install -r requirements-test.txt
           pip3 install --no-deps -e .
-          pip3 install git+https://github.com/ByteDance-Seed/VeOmni.git@v0.1.4
+          pip3 install git+https://github.com/ByteDance-Seed/VeOmni.git@v0.1.8 --ignore-requires-python --no-deps
+          pip3 install transformers==4.57.3
       - name: Prepare gsm8k dataset
         run: |
           ray stop --force

.github/workflows/e2e_sft_llm_ascend.yml

@@ -96,7 +96,8 @@ jobs:
       - name: Install the current repository
         run: |
           pip install --no-deps -e .
-          pip install git+https://github.com/ByteDance-Seed/VeOmni.git@v0.1.4
+          pip install git+https://github.com/ByteDance-Seed/VeOmni.git@v0.1.8 --ignore-requires-python --no-deps
+          pip install transformers==4.57.3
           pip install pandas==2.3.3
           pip uninstall -y mbridge
           pip install git+https://github.com/ISEEKYAN/mbridge.git@89eb10

.github/workflows/e2e_sft_vlm.yml

@@ -105,7 +105,8 @@ jobs:
           pip3 install peft
           pip3 install -r requirements-test.txt
           pip3 install --no-deps -e .
-          pip3 install git+https://github.com/ByteDance-Seed/VeOmni.git@v0.1.4
+          pip3 install git+https://github.com/ByteDance-Seed/VeOmni.git@v0.1.8 --ignore-requires-python --no-deps
+          pip3 install transformers==4.57.3
       - name: Prepare pokemon-gpt4o-captions dataset
         run: |
           ray stop --force

tests/special_e2e/sft/test_sft_engine_all.sh

@@ -25,7 +25,7 @@ echo "run with sp2 fsdp_size2 num_gpus8 fsdp_strategy fsdp2"
 BACKEND=fsdp SP_SIZE=2 FSDP_SIZE=2 NUM_GPUS=8 FSDP_STRATEGY=fsdp2 bash tests/special_e2e/sft/run_sft_engine.sh
 
 # test with veomni
-echo "run with sp2 fsdp_size4 num_gpus8 fsdp_strategy fsdp2"
+echo "run with sp2 fsdp_size4 num_gpus8 fsdp_strategy fsdp2 backend veomni"
 BACKEND=veomni SP_SIZE=2 FSDP_SIZE=4 NUM_GPUS=8 FSDP_STRATEGY=fsdp2 bash tests/special_e2e/sft/run_sft_engine.sh
 
 

verl/workers/config/engine.py

@@ -12,6 +12,8 @@
 # See the License for the specific language governing permissions and
 # limitations under the License.
 
+import logging
+import os
 import warnings
 from dataclasses import dataclass, field
 from typing import Any, Callable, Literal, Optional
@@ -37,6 +39,10 @@
 ]
 
 
+logger = logging.getLogger(__name__)
+logger.setLevel(os.getenv("VERL_LOGGING_LEVEL", "INFO"))
+
+
 # TODO: rename to RouterReplayConfig after removing the legacy implementation
 @dataclass
 class EngineRouterReplayConfig(BaseConfig):
@@ -311,6 +317,8 @@ class VeOmniEngineConfig(EngineConfig):
 
     """
 
+    _mutable_fields = EngineConfig._mutable_fields | {"attn_implementation"}
+
     wrap_policy: dict[str, Any] = field(default_factory=dict)
     offload_policy: bool = False
     reshard_after_forward: bool = True
@@ -342,6 +350,16 @@ def __post_init__(self):
         super().__post_init__()
         assert self.strategy in ["veomni"], f"strategy {self.strategy} not supported"
 
+        replacements = {
+            "flash_attention_2": "veomni_flash_attention_2_with_sp",
+            "flash_attention_3": "veomni_flash_attention_3_with_sp",
+            "flash_attention_4": "veomni_flash_attention_4_with_sp",
+        }
+        if self.attn_implementation in replacements:
+            new_impl = replacements[self.attn_implementation]
+            logger.info(f"Replacing attn_implementation from '{self.attn_implementation}' to '{new_impl}'")
+            self.attn_implementation = new_impl
+
 
 @dataclass
 class TorchtitanEngineConfig(EngineConfig):

verl/workers/engine/veomni/transformer_impl.py

@@ -26,6 +26,7 @@
 from veomni.distributed.torch_parallelize import build_parallelize_model
 from veomni.models.auto import build_foundation_model
 from veomni.optim import build_lr_scheduler, build_optimizer
+from veomni.utils.seqlen_pos_transform_utils import prepare_fa_kwargs_from_position_ids
 
 import verl.utils.torch_functional as verl_F
 from verl.trainer.config import CheckpointConfig
@@ -102,7 +103,7 @@ def __init__(
             dp_size=dp_size,
             dp_replicate_size=data_parallel_replicate_size,
             dp_shard_size=data_parallel_shard_size,
-            ep_size=self.engine_config.expert_parallel_size,
+            extra_parallel_sizes=(self.engine_config.expert_parallel_size,),
             ulysses_size=self.engine_config.ulysses_parallel_size,
             dp_mode=self.data_parallel_mode,
         )
@@ -214,7 +215,9 @@ def _build_model_optimizer(self):
             enable_mixed_precision=self.engine_config.mixed_precision,
             enable_gradient_checkpointing=self.model_config.enable_gradient_checkpointing,
             enable_fsdp_offload=self.engine_config.enable_fsdp_offload,
-            basic_modules=module._no_split_modules + self.engine_config.basic_modules,
+            basic_modules=list(
+                set(getattr(module, "_no_split_modules", None) or []) | set(self.engine_config.basic_modules)
+            ),
             enable_reentrant=self.engine_config.enable_reentrant,
             enable_forward_prefetch=self.engine_config.forward_prefetch,
         )
@@ -575,19 +578,60 @@ def __call__(self, batch: Sequence[dict[str, "torch.Tensor"]]) -> dict[str, "tor
         return batch
 
 
+def _prepare_veomni_flash_attention_kwargs(position_ids: torch.Tensor) -> dict[str, torch.Tensor | int]:
+    """Normalize packed position_ids layout and derive varlen FlashAttention kwargs.
+
+    Supported formats for use_remove_padding=true:
+        - 2D: (1, total_nnz) - standard packed format
+        - 3D: (rope_dim, 1, total_nnz) - VeRL mRoPE packed format
+    """
+    if position_ids.dim() == 2:
+        # (1, total_nnz) - standard packed format
+        fa_position_ids = position_ids
+    elif position_ids.dim() == 3:
+        # (rope_dim, 1, total_nnz) - VeRL mRoPE packed format
+        if position_ids.shape[1] == 1:
+            fa_position_ids = position_ids[0]
+        else:
+            raise ValueError(
+                f"Unsupported 3D position_ids shape: {tuple(position_ids.shape)}, expected (rope_dim, 1, total_nnz)"
+            )
+    else:
+        raise ValueError(
+            f"Unsupported position_ids rank: {position_ids.dim()}, "
+            f"expected 2 (1, total_nnz) or 3 (rope_dim, 1, total_nnz)"
+        )
+
+    (cu_seq_lens_q, cu_seq_lens_k), (max_length_q, max_length_k) = prepare_fa_kwargs_from_position_ids(fa_position_ids)
+    return {
+        "cu_seq_lens_q": cu_seq_lens_q,
+        "cu_seq_lens_k": cu_seq_lens_k,
+        "max_length_q": max_length_q,
+        "max_length_k": max_length_k,
+    }
+
+
 @EngineRegistry.register(model_type="language_model", backend=["veomni"], device=["cuda", "npu"])
 class VeOmniEngineWithLMHead(VeOmniEngine, FSDPEngineWithLMHead):
     def prepare_model_inputs(self, micro_batch: TensorDict):
         # TODO: Cannot work properly for qwen_vl ulysses
         model_inputs, output_args = super().prepare_model_inputs(micro_batch)
         input_ids_rmpad = model_inputs["input_ids"]
+        sp_enabled = parallel_state.get_parallel_state().sp_enabled
+        sp_shard_collator = OmniSequenceShardCollator() if sp_enabled else None
+
         if self.module.config.model_type in VL_TYPE2INDEX.keys():
             image_mask = input_ids_rmpad == VL_TYPE2INDEX[self.module.config.model_type]["IMAGE_INPUT_INDEX"]
             video_mask = input_ids_rmpad == VL_TYPE2INDEX[self.module.config.model_type]["VIDEO_INPUT_INDEX"]
             model_inputs.update({"image_mask": image_mask, "video_mask": video_mask})
 
-            if parallel_state.get_parallel_state().sp_enabled:
-                omni_sequence_shard_collator = OmniSequenceShardCollator()
-                omni_sequence_shard_collator(model_inputs)
+            if sp_enabled:
+                sp_shard_collator(model_inputs)
+
+        use_remove_padding = tu.get_non_tensor_data(data=micro_batch, key="use_remove_padding", default=True)
+        if use_remove_padding and model_inputs.get("position_ids", None) is not None:
+            model_inputs.update(_prepare_veomni_flash_attention_kwargs(model_inputs["position_ids"]))
+            if sp_enabled:
+                model_inputs["position_ids"] = sp_shard_collator.sp_slice(model_inputs["position_ids"], dim=-1)
 
         return model_inputs, output_args