[trainer] feat: VeOmniEngine supports qwen3_vl ulysses (verl-project#4806)

### What does this PR do?

as title.

### Checklist Before Starting

- [x] Search for similar PRs. Paste at least one query link here: ...
- [x] Format the PR title as `[{modules}] {type}: {description}` (This
will be checked by the CI)
- `{modules}` include `fsdp`, `megatron`, `sglang`, `vllm`, `rollout`,
`trainer`, `ci`, `training_utils`, `recipe`, `hardware`, `deployment`,
`ray`, `worker`, `single_controller`, `misc`, `perf`, `model`, `algo`,
`env`, `tool`, `ckpt`, `doc`, `data`, `cfg`, `reward`
- If this PR involves multiple modules, separate them with `,` like
`[megatron, fsdp, doc]`
  - `{type}` is in `feat`, `fix`, `refactor`, `chore`, `test`
- If this PR breaks any API (CLI arguments, config, function signature,
etc.), add `[BREAKING]` to the beginning of the title.
  - Example: `[BREAKING][fsdp, megatron] feat: dynamic batching`

### Test

> For changes that can not be tested by CI (e.g., algorithm
implementation, new model support), validate by experiment(s) and show
results like training curve plots, evaluation results, etc.

### API and Usage Example

> Demonstrate how the API changes if any, and provide usage example(s)
if possible.

```python
# Add code snippet or script demonstrating how to use this
```

### Design & Code Changes

> Demonstrate the high-level design if this PR is complex, and list the
specific changes.

### Checklist Before Submitting

> [!IMPORTANT]
> Please check all the following items before requesting a review,
otherwise the reviewer might deprioritize this PR for review.

- [x] Read the [Contribute
Guide](https://github.com/volcengine/verl/blob/main/CONTRIBUTING.md).
- [x] Apply [pre-commit
checks](https://github.com/volcengine/verl/blob/main/CONTRIBUTING.md#code-linting-and-formatting):
`pre-commit install && pre-commit run --all-files --show-diff-on-failure
--color=always`
- [ ] Add / Update [the
documentation](https://github.com/volcengine/verl/tree/main/docs).
- [ ] Add unit or end-to-end test(s) to [the CI
workflow](https://github.com/volcengine/verl/tree/main/.github/workflows)
to cover all the code. If not feasible, explain why: ...
- [ ] Once your PR is ready for CI, send a message in [the `ci-request`
channel](https://verl-project.slack.com/archives/C091TCESWB1) in [the
`verl` Slack
workspace](https://join.slack.com/t/verl-project/shared_invite/zt-3855yhg8g-CTkqXu~hKojPCmo7k_yXTQ).
(If not accessible, please try [the Feishu group
(飞书群)](https://applink.larkoffice.com/client/chat/chatter/add_by_link?link_token=772jd4f1-cd91-441e-a820-498c6614126a).)

Author: A1waysBeenHere

tests/special_e2e/sft/test_sft_engine_all.sh

@@ -25,9 +25,8 @@ 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
-# FIXME(ji-huazhong): set SP=1 cause qwen_vl do not support SP right now
-echo "run with sp1 fsdp_size4 num_gpus8 fsdp_strategy fsdp2"
-BACKEND=veomni SP_SIZE=1 FSDP_SIZE=8 NUM_GPUS=8 FSDP_STRATEGY=fsdp2 bash tests/special_e2e/sft/run_sft_engine.sh
+echo "run with sp2 fsdp_size4 num_gpus8 fsdp_strategy fsdp2"
+BACKEND=veomni SP_SIZE=2 FSDP_SIZE=4 NUM_GPUS=8 FSDP_STRATEGY=fsdp2 bash tests/special_e2e/sft/run_sft_engine.sh
 
 
 # test with megatron

verl/workers/engine/veomni/transformer_impl.py

@@ -14,7 +14,8 @@
 
 
 import logging
-from typing import Any, Callable
+from dataclasses import dataclass, field
+from typing import Any, Callable, Sequence
 
 import torch
 import torch.distributed as dist
@@ -133,7 +134,6 @@ def initialize(self):
             attn_implementation=self.engine_config.attn_implementation,
             moe_implementation=self.engine_config.moe_implementation,
             init_device=self.engine_config.init_device,
-            force_use_huggingface=self.engine_config.force_use_huggingface,
         )
 
         module_config = self.module.config
@@ -250,10 +250,7 @@ def forward_backward_batch(self, data: TensorDict, loss_function: Callable, forw
         return postprocess_batch_func(output_lst=output_lst, indices=indices, data=data)
 
     def get_data_parallel_rank(self):
-        if parallel_state.get_parallel_state().ulysses_size > 1:
-            return parallel_state.get_parallel_state().device_mesh["dp"].get_local_rank()
-        else:
-            return torch.distributed.get_rank()
+        return parallel_state.get_parallel_state().device_mesh.get_local_rank("dp")
 
     def get_data_parallel_size(self):
         return torch.distributed.get_world_size() // parallel_state.get_parallel_state().ulysses_size
@@ -299,7 +296,7 @@ def __enter__(self):
         assert isinstance(self.engine, VeOmniEngine)
         super().__enter__()
         self.engine.ulysses_sharding_manager.__enter__()
-        self.engine.module.eval()
+        self.engine.module.train()
 
     def __exit__(self, exc_type, exc_value, traceback):
         assert isinstance(self.engine, VeOmniEngine)
@@ -333,6 +330,41 @@ def __exit__(self, exc_type, exc_value, traceback):
         super().__exit__(exc_type, exc_value, traceback)
 
 
+@dataclass
+class OmniSequenceShardCollator:
+    """
+    Data collator to chunk inputs along the sequence length.
+    """
+
+    # features to slice sequence dimension
+    sp_slice_features: dict[str, int] = field(
+        default_factory=lambda: {
+            "input_ids": -1,
+            "labels": -1,
+            "pixel_values": 0,
+            "pixel_values_videos": 0,
+        },
+        metadata={"help": "features to slice sequence dimension."},
+    )
+
+    def __post_init__(self):
+        self.sp_size = parallel_state.get_parallel_state().sp_size
+        self.sp_rank = parallel_state.get_parallel_state().sp_rank
+
+    def sp_slice(self, feature: torch.Tensor, dim: int = -1) -> dict[str, "torch.Tensor"]:
+        seq_length = feature.size(dim)
+        sp_chunk_size = (seq_length + self.sp_size - 1) // self.sp_size
+        return feature.narrow(dim, self.sp_rank * sp_chunk_size, sp_chunk_size)
+
+    def __call__(self, batch: Sequence[dict[str, "torch.Tensor"]]) -> dict[str, "torch.Tensor"]:
+        # sp slice
+        for key in batch.keys():
+            if key in self.sp_slice_features.keys():
+                batch[key] = self.sp_slice(batch[key], dim=self.sp_slice_features[key])
+
+        return batch
+
+
 @EngineRegistry.register(model_type="language_model", backend=["veomni"], device=["cuda", "npu"])
 class VeOmniEngineWithLMHead(VeOmniEngine, FSDPEngineWithLMHead):
     def prepare_model_inputs(self, micro_batch: TensorDict):
@@ -344,4 +376,8 @@ def prepare_model_inputs(self, micro_batch: TensorDict):
             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)
+
         return model_inputs, output_args