Support Qwen3.5

.github/workflows/docker/docker-compose.yaml

@@ -1,6 +1,6 @@
 services:
   trinity-node-1:
-    image: trinity-rft-unittest:20260228
+    image: trinity-rft-unittest:20260310
     cap_add:
       - SYS_PTRACE
     pull_policy: never
@@ -34,7 +34,7 @@ services:
             capabilities: [gpu]
 
   trinity-node-2:
-    image: trinity-rft-unittest:20260228
+    image: trinity-rft-unittest:20260310
     cap_add:
       - SYS_PTRACE
     pull_policy: never

pyproject.toml

@@ -52,11 +52,12 @@ trinity = "trinity.cli.launcher:main"
 
 [project.optional-dependencies]
 vllm = [
-    "vllm>=0.10.2,<=0.16.0,!=0.11.0,!=0.12.0",
+    "vllm>=0.10.2,<=0.17.1,!=0.11.0,!=0.12.0",
     # v0.11 has bug when prefix-caching is enabled so we exclude it
     # v0.12 has a huge performance regression so we exclude it
-    # v0.10.2 is the most stable version, but we allow up to 0.16.0 for new features
-    # v0.16.0 is required for transformers>=5.0.0
+    # v0.10.2 is the most stable version, but we allow up to 0.17.1 for new features
+    # For v0.16 and v0.17, the default dependencies require transformers < 5.
+    # We have patched vLLM to support transformers >= 5.0.0.
 ]
 data = [
     "py-data-juicer>=1.4.3"
@@ -80,10 +81,10 @@ dev = [
     "viztracer",
 ]
 megatron = [
-    "megatron-core[mlm]==0.15.0",
+    "megatron-core[mlm]>=0.15.0",
     # if you found "undefined symbol" error in transformer engine
     # reinstall it with --no-build-isolation and `--no-cache-dir` flag
-    # "transformer_engine[pytorch]==2.10.0",
+    "transformer_engine[pytorch]>=2.10.0",
 
     # Install mbridge from main branch (unreleased version)
     # "mbridge @ git+https://github.com/ISEEKYAN/mbridge.git@20e9ffbbe72ae7b1df83bfe1bc3c11f7382f2612",
@@ -109,7 +110,7 @@ mm = [
 ]
 
 flash_attn = [
-    "flash-attn==2.8.1"
+    "flash-attn>=2.8.1"
 ]
 
 [tool.setuptools.packages.find]

trinity/common/models/vllm_model.py

@@ -63,21 +63,24 @@ def __init__(
             os.environ["VLLM_CACHE_ROOT"] = os.path.expanduser(
                 f"~/.cache/vllm/{config.bundle_indices}"
             )
+        self.tokenization_kwargs = {
+            "truncate_prompt_tokens": config.max_prompt_tokens
+            if config.enable_prompt_truncation
+            else None
+        }
         self.default_sampling_params = vllm.SamplingParams(
             n=1,
             temperature=config.temperature,
             max_tokens=config.max_response_tokens,
             min_tokens=config.min_response_tokens,
-            truncate_prompt_tokens=(
-                config.max_prompt_tokens if config.enable_prompt_truncation else None
-            ),
             skip_special_tokens=True,
             include_stop_str_in_output=False,
             output_kind=RequestOutputKind.FINAL_ONLY,
             logprobs=config.logprobs,
             top_p=config.top_p,
             top_k=config.top_k,
             ignore_eos=config.ignore_eos,
+            **(self.tokenization_kwargs if self.vllm_version <= parse_version("0.16.0") else {}),
         )
         self.ray_namespace = config.ray_namespace
         self.request_id = 0
@@ -417,11 +420,17 @@ async def sample(
     async def _generate_internal(self, prompt: Any, lora_request=None, **kwargs) -> Any:
         # Send the request to the LLM engine.
         self.request_id += 1
+        generate_kwargs = (
+            {"tokenization_kwargs": self.tokenization_kwargs}
+            if self.vllm_version > parse_version("0.16.0")
+            else {}
+        )
         stream = self.async_llm.generate(
             request_id=str(self.request_id),
             prompt=prompt,
             sampling_params=self._create_sampling_params(**kwargs),
             lora_request=lora_request,
+            **generate_kwargs,
         )
 
         # Consume the stream until the request is finished.

trinity/common/models/vllm_patch/__init__.py

@@ -17,8 +17,22 @@ def vllm_patch():
 
     trf_version = parse_version(transformers.__version__)
     vllm_version = parse_version(vllm.__version__)
-    if trf_version >= parse_version("5.0.0") and vllm_version < parse_version("0.16.0"):
-        raise ImportError("Please upgrade vllm to 0.16.0 or above to use transformers>=5.0.0.")
+    if trf_version >= parse_version("5.0.0"):
+        if vllm_version < parse_version("0.16.0"):
+            raise ImportError("Please upgrade vllm to 0.16.0 or above to use transformers>=5.0.0.")
+
+        from transformers.configuration_utils import PreTrainedConfig
+
+        original_init = PreTrainedConfig.__init__
+
+        def new_init(self, *args, **kwargs):
+            if "ignore_keys_at_rope_validation" in kwargs:
+                kwargs["ignore_keys_at_rope_validation"] = set(
+                    kwargs["ignore_keys_at_rope_validation"]
+                )
+            original_init(self, *args, **kwargs)
+
+        PreTrainedConfig.__init__ = new_init
 
 
 def get_vllm_version():

trinity/common/models/vllm_patch/worker_patch.py

@@ -13,10 +13,10 @@
 def patch_vllm_prompt_logprobs(model_runner: GPUModelRunner):  # noqa: C901
     """Patch vLLM model runner to support prompt logprobs extraction."""
     version = get_vllm_version()
-    if version < parse_version("0.10.2") or version > parse_version("0.16.0"):
+    if version < parse_version("0.10.2") or version >= parse_version("0.18.0"):
         raise ValueError(
             f"Unsupported vllm version: {vllm.__version__}. "
-            "This patch requires vllm version >= 0.10.2, <= 0.16.0."
+            "This patch requires vllm version >= 0.10.2, < 0.18.0."
         )
     is_v0102 = version == parse_version("0.10.2")
 

trinity/common/patch/qwen3_5.py

@@ -0,0 +1,201 @@
+from dataclasses import dataclass
+from functools import wraps
+from typing import Optional
+
+import torch
+from transformers.models.qwen3_5.modeling_qwen3_5 import (
+    BaseModelOutputWithPast,
+    Cache,
+    Qwen3_5CausalLMOutputWithPast,
+    Qwen3_5DynamicCache,
+    Qwen3_5ForConditionalGeneration,
+    Qwen3_5ModelOutputWithPast,
+    TransformersKwargs,
+    Unpack,
+    capture_outputs,
+    create_causal_mask,
+    merge_with_config_defaults,
+)
+
+
+# TODO: may optimize this function
+def ulysses_gated_delta_net_forward_decorator(func):
+    @wraps(func)
+    def wrapper(
+        hidden_states: torch.Tensor,
+        cache_params: Qwen3_5DynamicCache | None = None,
+        cache_position: torch.LongTensor | None = None,
+        attention_mask: torch.Tensor | None = None,
+    ):
+        from verl.utils.ulysses import (
+            gather_outputs_and_unpad,
+            get_ulysses_sequence_parallel_world_size,
+            slice_input_tensor,
+        )
+
+        ulysses_sp_size = get_ulysses_sequence_parallel_world_size()
+        if ulysses_sp_size > 1:
+            hidden_states = gather_outputs_and_unpad(hidden_states, gather_dim=1)
+
+        output = func(hidden_states, cache_params, cache_position, attention_mask)
+
+        if ulysses_sp_size > 1:
+            output = slice_input_tensor(output, dim=1, padding=False)
+        return output
+
+    return wrapper
+
+
+@merge_with_config_defaults
+@capture_outputs
+def qwen35_text_forward(
+    self,
+    input_ids: torch.LongTensor | None = None,
+    attention_mask: torch.Tensor | None = None,
+    position_ids: torch.LongTensor | None = None,
+    past_key_values: Cache | None = None,
+    inputs_embeds: torch.FloatTensor | None = None,
+    use_cache: bool | None = None,
+    cache_position: torch.LongTensor | None = None,
+    **kwargs: Unpack[TransformersKwargs],
+) -> BaseModelOutputWithPast:
+    if (input_ids is None) ^ (inputs_embeds is not None):
+        raise ValueError("You must specify exactly one of input_ids or inputs_embeds")
+
+    if inputs_embeds is None:
+        inputs_embeds = self.embed_tokens(input_ids)
+
+    if use_cache and past_key_values is None:
+        past_key_values = Qwen3_5DynamicCache(config=self.config)
+
+    if cache_position is None:
+        past_seen_tokens = past_key_values.get_seq_length() if past_key_values is not None else 0
+        cache_position = torch.arange(
+            past_seen_tokens, past_seen_tokens + inputs_embeds.shape[1], device=inputs_embeds.device
+        )
+
+    # mrope: the hard coded `3` is for temporal, height and width.
+    if position_ids is None:
+        position_ids = cache_position.view(1, 1, -1).expand(3, inputs_embeds.shape[0], -1)
+    elif position_ids.ndim == 2:
+        position_ids = position_ids[None, ...].expand(3, position_ids.shape[0], -1)
+
+    if position_ids.ndim == 3 and position_ids.shape[0] == 4:
+        text_position_ids = position_ids[0]
+        position_ids = position_ids[1:]
+    else:
+        text_position_ids = position_ids[0]
+
+    causal_mask = create_causal_mask(
+        config=self.config,
+        inputs_embeds=inputs_embeds,
+        attention_mask=attention_mask,
+        cache_position=cache_position,
+        past_key_values=past_key_values,
+        position_ids=text_position_ids,
+    )
+    linear_attn_mask = self._update_linear_attn_mask(attention_mask, cache_position)
+
+    hidden_states = inputs_embeds
+    position_embeddings = self.rotary_emb(hidden_states, position_ids)
+
+    for layer_idx, decoder_layer in enumerate(self.layers[: self.config.num_hidden_layers]):
+        layer_mask = (
+            linear_attn_mask if decoder_layer.layer_type == "linear_attention" else causal_mask
+        )
+
+        hidden_states = decoder_layer(
+            hidden_states,
+            position_embeddings=position_embeddings,
+            attention_mask=layer_mask,
+            position_ids=text_position_ids,
+            past_key_values=past_key_values,
+            use_cache=use_cache,
+            cache_position=cache_position,
+            **kwargs,
+        )
+
+    hidden_states = self.norm(hidden_states)
+
+    return Qwen3_5ModelOutputWithPast(
+        last_hidden_state=hidden_states,
+        past_key_values=past_key_values,
+    )
+
+
+@dataclass
+class Qwen3_5CausalLMOutputForPPO(Qwen3_5CausalLMOutputWithPast):
+    log_probs: Optional[torch.FloatTensor] = None
+    entropy: Optional[torch.FloatTensor] = None
+
+
+def forward_with_torch_backend(
+    self: Qwen3_5ForConditionalGeneration,
+    input_ids: torch.LongTensor = None,
+    labels: Optional[torch.LongTensor] = None,
+    temperature: float = 1.0,
+    **kwargs,
+) -> tuple | Qwen3_5CausalLMOutputForPPO:
+    from verl.utils.experimental.torch_functional import FusedLinearForPPO
+
+    outputs = self.model(input_ids=input_ids, **kwargs)
+    hidden_states = outputs[0]
+
+    # Loss calculations
+    if labels is not None:
+        rolled_labels = torch.roll(labels, shifts=-1, dims=-1)
+    elif input_ids is not None:
+        rolled_labels = torch.roll(input_ids, shifts=-1, dims=-1)
+    else:
+        raise RuntimeError(
+            "To use forward_with_torch_backend, either labels or input_ids must be provided."
+        )
+
+    fused_linear_for_ppo = FusedLinearForPPO()
+    log_probs, entropy = fused_linear_for_ppo.forward(
+        hidden_states=hidden_states,
+        vocab_weights=self.lm_head.weight,
+        input_ids=rolled_labels,
+        temperature=temperature,
+    )
+    return Qwen3_5CausalLMOutputForPPO(
+        log_probs=log_probs,
+        entropy=entropy,
+        hidden_states=outputs.hidden_states,
+    )
+
+
+def forward_with_triton_backend(
+    self: Qwen3_5ForConditionalGeneration,
+    input_ids: torch.LongTensor = None,
+    labels: Optional[torch.LongTensor] = None,
+    temperature: float = 1.0,
+    **kwargs,
+) -> tuple | Qwen3_5CausalLMOutputForPPO:
+    from verl.utils.kernel.linear_cross_entropy import linear_cross_entropy
+
+    outputs = self.model(input_ids=input_ids, **kwargs)
+    hidden_states = outputs[0]
+
+    # Loss calculations
+    if labels is not None:
+        rolled_labels = torch.roll(labels, shifts=-1, dims=-1)
+    elif input_ids is not None:
+        rolled_labels = torch.roll(input_ids, shifts=-1, dims=-1)
+    else:
+        raise RuntimeError(
+            "To use forward_with_triton_backend, either labels or input_ids must be provided."
+        )
+
+    log_probs, entropy = linear_cross_entropy(
+        hidden_states,
+        self.lm_head.weight,
+        rolled_labels,
+        temperature,
+        "none",
+    )
+    return Qwen3_5CausalLMOutputForPPO(
+        log_probs=log_probs,
+        entropy=entropy,
+        hidden_states=outputs.hidden_states,
+    )