From a93566dcf8f961f6327c8a1704d3f97bd45d0386 Mon Sep 17 00:00:00 2001 From: ys2025-AI Date: Mon, 31 Aug 2026 00:28:10 +0000 Subject: [PATCH 1/2] [model] refactor qwen3.5 GDN to use FLA causal conv1d and cu_seq_lens_q - Add FLA causal_conv1d as the conv function for packed sequences - Pass cu_seq_lens_q directly to forward, removing manual padding/unpadding - Derive layer_types from linear_attention_freq in __init__ - Use layer_number-1 to match HF indexing --- src/mcore_bridge/model/mm_gpts/qwen3_5.py | 51 +++++++++++++---------- 1 file changed, 30 insertions(+), 21 deletions(-) diff --git a/src/mcore_bridge/model/mm_gpts/qwen3_5.py b/src/mcore_bridge/model/mm_gpts/qwen3_5.py index b673bf7..dd843a3 100644 --- a/src/mcore_bridge/model/mm_gpts/qwen3_5.py +++ b/src/mcore_bridge/model/mm_gpts/qwen3_5.py @@ -1,5 +1,7 @@ # Copyright (c) ModelScope Contributors. All rights reserved. import torch +import torch.nn.functional as F +from fla.modules.conv.causal_conv1d import causal_conv1d as _fla_causal_conv1d from megatron.core.extensions.transformer_engine import _get_extra_te_kwargs from megatron.core.models.huggingface import HuggingFaceModule as _HuggingFaceModule from megatron.core.tensor_parallel import gather_from_sequence_parallel_region, scatter_to_sequence_parallel_region @@ -25,41 +27,48 @@ class Qwen3_5MoeGatedDeltaNet(_HuggingFaceModule, _Qwen3_5MoeGatedDeltaNet): def __init__(self, config: TransformerConfig, submodules: SelfAttentionSubmodules, layer_number: int, **kwargs): assert config.context_parallel_size == 1, 'Qwen3_5 currently does not support context parallel.' assert _Qwen3_5MoeGatedDeltaNet is not object, 'please update the `transformers` version.' - _Qwen3_5MoeGatedDeltaNet.__init__(self, config, layer_number) + if getattr(config, 'layer_types', None) is None: + freq = config.linear_attention_freq + if isinstance(freq, int): + freq = [0 if (i + 1) % freq == 0 else 1 for i in range(config.num_layers)] + config.layer_types = ['linear_attention' if f else 'full_attention' for f in freq] + _Qwen3_5MoeGatedDeltaNet.__init__(self, config, layer_number - 1) self.config = config extra_kwargs = _get_extra_te_kwargs(config) self.to(dtype=extra_kwargs['params_dtype'], device=extra_kwargs['device']) + self.causal_conv1d_fn = self._fla_conv_fn + self._cur_cu_seqlens = None + + def _fla_conv_fn(self, x, weight, bias, activation, seq_idx=None): + x_btd = x.transpose(1, 2) + if isinstance(activation, str): + act = activation + elif activation is F.silu or activation == 'silu': + act = 'silu' + else: + act = None + out, _ = _fla_causal_conv1d(x_btd, weight=weight, bias=bias, activation=act, cu_seqlens=self._cur_cu_seqlens) + return out.transpose(1, 2) def forward(self, hidden_states: torch.Tensor, **kwargs): config = self.config if config.sequence_parallel and config.tensor_model_parallel_size > 1: hidden_states = gather_from_sequence_parallel_region(hidden_states, tensor_parallel_output_grad=False) - seq_len = hidden_states.shape[0] packed_seq_params = kwargs.get('packed_seq_params') thd_format = packed_seq_params is not None and packed_seq_params.qkv_format == 'thd' - # Note: for packed inputs, we do not perform padding_free unpadding. - # Doing so would allow different sequences to see each other; for efficiency we keep this implementation. if thd_format: - max_seqlen_q = int(packed_seq_params.max_seqlen_q) - new_hidden_states = hidden_states.new_zeros( - (packed_seq_params.num_samples, max_seqlen_q, hidden_states.shape[-1])) - attention_mask = hidden_states.new_zeros((packed_seq_params.num_samples, max_seqlen_q), dtype=torch.bool) - cu_seqlens_q = packed_seq_params.cu_seqlens_q - for i in range(packed_seq_params.num_samples): - start, end = cu_seqlens_q[i], cu_seqlens_q[i + 1] - attention_mask[i, :end - start] = True - new_hidden_states[i, :end - start] = hidden_states[start:end, 0] - hidden_states = new_hidden_states + cu_seq_lens_q = packed_seq_params.cu_seqlens_q + self._cur_cu_seqlens = cu_seq_lens_q + attention_mask = None else: - hidden_states = hidden_states.transpose(0, 1) + cu_seq_lens_q = None + self._cur_cu_seqlens = None attention_mask = resolve_gdn_attention_mask(kwargs) + hidden_states = hidden_states.transpose(0, 1) with get_cuda_rng_tracker().fork('data-parallel-rng'): - res = super().forward(hidden_states=hidden_states, attention_mask=attention_mask) - if thd_format: - res = res[attention_mask][:, None] - res = torch.concat([res, res.new_zeros(seq_len - res.shape[0], 1, res.shape[2])]) - else: - res = res.transpose(0, 1).contiguous() + res = super().forward( + hidden_states=hidden_states, attention_mask=attention_mask, cu_seq_lens_q=cu_seq_lens_q) + res = res.transpose(0, 1).contiguous() if config.sequence_parallel and config.tensor_model_parallel_size > 1: res = scatter_to_sequence_parallel_region(res) return res, None From 28893645053246d12935f176c1d4784c1103f543 Mon Sep 17 00:00:00 2001 From: ys2025-AI Date: Mon, 31 Aug 2026 10:47:38 +0000 Subject: [PATCH 2/2] [model] make qwen3.5 GDN FLA import optional and compatible across transformers versions --- src/mcore_bridge/model/mm_gpts/qwen3_5.py | 30 ++++++++++++++++++++--- 1 file changed, 26 insertions(+), 4 deletions(-) diff --git a/src/mcore_bridge/model/mm_gpts/qwen3_5.py b/src/mcore_bridge/model/mm_gpts/qwen3_5.py index dd843a3..f76ad95 100644 --- a/src/mcore_bridge/model/mm_gpts/qwen3_5.py +++ b/src/mcore_bridge/model/mm_gpts/qwen3_5.py @@ -1,13 +1,14 @@ # Copyright (c) ModelScope Contributors. All rights reserved. import torch import torch.nn.functional as F -from fla.modules.conv.causal_conv1d import causal_conv1d as _fla_causal_conv1d +import transformers from megatron.core.extensions.transformer_engine import _get_extra_te_kwargs from megatron.core.models.huggingface import HuggingFaceModule as _HuggingFaceModule from megatron.core.tensor_parallel import gather_from_sequence_parallel_region, scatter_to_sequence_parallel_region from megatron.core.tensor_parallel.random import get_cuda_rng_tracker from megatron.core.transformer.attention import SelfAttentionSubmodules from megatron.core.transformer.transformer_config import TransformerConfig +from packaging import version from mcore_bridge.utils import get_env_args @@ -21,6 +22,14 @@ except ImportError: _Qwen3_5MoeGatedDeltaNet = object +try: + from fla.modules.conv.causal_conv1d import causal_conv1d as _fla_causal_conv1d +except ImportError: + _fla_causal_conv1d = None + +_TRANSFORMERS_SUPPORTS_CU_SEQ_LENS_Q = version.parse(transformers.__version__) >= version.parse('5.9.0') +_TRANSFORMERS_HAS_CONV1D_FN = version.parse(transformers.__version__) < version.parse('5.15.0') + class Qwen3_5MoeGatedDeltaNet(_HuggingFaceModule, _Qwen3_5MoeGatedDeltaNet): @@ -36,7 +45,8 @@ def __init__(self, config: TransformerConfig, submodules: SelfAttentionSubmodule self.config = config extra_kwargs = _get_extra_te_kwargs(config) self.to(dtype=extra_kwargs['params_dtype'], device=extra_kwargs['device']) - self.causal_conv1d_fn = self._fla_conv_fn + if _fla_causal_conv1d is not None and _TRANSFORMERS_HAS_CONV1D_FN: + self.causal_conv1d_fn = self._fla_conv_fn self._cur_cu_seqlens = None def _fla_conv_fn(self, x, weight, bias, activation, seq_idx=None): @@ -57,6 +67,15 @@ def forward(self, hidden_states: torch.Tensor, **kwargs): packed_seq_params = kwargs.get('packed_seq_params') thd_format = packed_seq_params is not None and packed_seq_params.qkv_format == 'thd' if thd_format: + if not _TRANSFORMERS_SUPPORTS_CU_SEQ_LENS_Q: + raise RuntimeError( + 'Packed (thd) varlen forward for Qwen3.5 requires transformers>=5.9.0, since cu_seq_lens_q was ' + 'introduced in huggingface/transformers#45034. ' + f'Current transformers version is {transformers.__version__}.') + if _fla_causal_conv1d is None: + raise ImportError( + 'flash-linear-attention is required for packed (thd) varlen forward of Qwen3.5. ' + "Please install it via: `pip install -U 'flash-linear-attention' --no-build-isolation`") cu_seq_lens_q = packed_seq_params.cu_seqlens_q self._cur_cu_seqlens = cu_seq_lens_q attention_mask = None @@ -66,8 +85,11 @@ def forward(self, hidden_states: torch.Tensor, **kwargs): attention_mask = resolve_gdn_attention_mask(kwargs) hidden_states = hidden_states.transpose(0, 1) with get_cuda_rng_tracker().fork('data-parallel-rng'): - res = super().forward( - hidden_states=hidden_states, attention_mask=attention_mask, cu_seq_lens_q=cu_seq_lens_q) + if cu_seq_lens_q is not None: + res = super().forward( + hidden_states=hidden_states, attention_mask=attention_mask, cu_seq_lens_q=cu_seq_lens_q) + else: + res = super().forward(hidden_states=hidden_states, attention_mask=attention_mask) res = res.transpose(0, 1).contiguous() if config.sequence_parallel and config.tensor_model_parallel_size > 1: res = scatter_to_sequence_parallel_region(res)