feat: standardize import formatting and fix attention implementation string

- Change double quotes to single quotes for consistency in `attn_implementation` parameter
- Reformat multi-line imports to single line for better readability
- Remove unnecessary import error message in linear attention validation
- Maintain code style consistency across the codebase

Author: meichangsu1

cookbook/transformers/sp_fsdp_dense.py

@@ -70,7 +70,7 @@ def train():
         model_cls=TwinkleQwen3_5ForCausalLM,
         device_mesh=device_mesh,
         strategy='native_fsdp',
-        attn_implementation="flash_attention_2"
+        attn_implementation='flash_attention_2'
     )
 
     lora_config = LoraConfig(target_modules='all-linear', lora_dropout=0.0)

src/twinkle/model/transformers/__init__.py

@@ -4,13 +4,8 @@
 from twinkle.utils.import_utils import _LazyModule
 
 if TYPE_CHECKING:
-    from .models import (
-        TwinkleQwen3_5DecoderLayer,
-        TwinkleQwen3_5ForCausalLM,
-        TwinkleQwen3_5GatedDeltaNet,
-        TwinkleQwen3_5PreTrainedModel,
-        TwinkleQwen3_5TextModel,
-    )
+    from .models import (TwinkleQwen3_5DecoderLayer, TwinkleQwen3_5ForCausalLM, TwinkleQwen3_5GatedDeltaNet,
+                         TwinkleQwen3_5PreTrainedModel, TwinkleQwen3_5TextModel)
     from .multi_lora_transformers import MultiLoraTransformersModel
     from .transformers import TransformersModel
 else:

src/twinkle/model/transformers/models/__init__.py

@@ -1,11 +1,6 @@
 # Copyright (c) ModelScope Contributors. All rights reserved.
-from .qwen3_5 import (
-    TwinkleQwen3_5DecoderLayer,
-    TwinkleQwen3_5ForCausalLM,
-    TwinkleQwen3_5GatedDeltaNet,
-    TwinkleQwen3_5PreTrainedModel,
-    TwinkleQwen3_5TextModel,
-)
+from .qwen3_5 import (TwinkleQwen3_5DecoderLayer, TwinkleQwen3_5ForCausalLM, TwinkleQwen3_5GatedDeltaNet,
+                      TwinkleQwen3_5PreTrainedModel, TwinkleQwen3_5TextModel)
 
 __all__ = [
     'TwinkleQwen3_5PreTrainedModel',

src/twinkle/model/transformers/models/qwen3_5/__init__.py

@@ -1,11 +1,6 @@
 # Copyright (c) ModelScope Contributors. All rights reserved.
-from .modeling_qwen3_5 import (
-    TwinkleQwen3_5DecoderLayer,
-    TwinkleQwen3_5ForCausalLM,
-    TwinkleQwen3_5GatedDeltaNet,
-    TwinkleQwen3_5PreTrainedModel,
-    TwinkleQwen3_5TextModel,
-)
+from .modeling_qwen3_5 import (TwinkleQwen3_5DecoderLayer, TwinkleQwen3_5ForCausalLM, TwinkleQwen3_5GatedDeltaNet,
+                               TwinkleQwen3_5PreTrainedModel, TwinkleQwen3_5TextModel)
 
 __all__ = [
     'TwinkleQwen3_5PreTrainedModel',

src/twinkle/model/transformers/models/qwen3_5/modeling_qwen3_5.py

@@ -2,21 +2,19 @@
 from __future__ import annotations
 
 import importlib.util
-from typing import Any, Callable, Optional
-
 import torch
 import torch.nn.functional as F
 from torch import nn
 from transformers.cache_utils import Cache
 from transformers.generation import GenerationMixin
 from transformers.modeling_outputs import BaseModelOutputWithPast, CausalLMOutputWithPast
-from transformers.models.qwen3_5.configuration_qwen3_5 import Qwen3_5Config, Qwen3_5TextConfig
 from transformers.models.qwen3_5 import modeling_qwen3_5 as hf_qwen35
+from transformers.models.qwen3_5.configuration_qwen3_5 import Qwen3_5Config, Qwen3_5TextConfig
 from transformers.processing_utils import Unpack
 from transformers.utils import TransformersKwargs, can_return_tuple
 from transformers.utils.generic import merge_with_config_defaults
 from transformers.utils.output_capturing import capture_outputs
-
+from typing import Any, Callable, Optional
 
 try:
     from fla.modules import FusedRMSNormGated as _FLA_FUSED_RMS_NORM_GATED
@@ -37,10 +35,8 @@
 def _ensure_text_config(config: Qwen3_5TextConfig) -> Qwen3_5TextConfig:
     if isinstance(config, Qwen3_5TextConfig):
         return config
-    raise TypeError(
-        'TwinkleQwen3_5 text-only models require transformers.models.qwen3_5.Qwen3_5TextConfig. '
-        f'Got {type(config).__name__}.'
-    )
+    raise TypeError('TwinkleQwen3_5 text-only models require transformers.models.qwen3_5.Qwen3_5TextConfig. '
+                    f'Got {type(config).__name__}.')
 
 
 def _ensure_linear_attention_fast_path() -> None:
@@ -52,10 +48,8 @@ def _ensure_linear_attention_fast_path() -> None:
     if not _HAS_CAUSAL_CONV1D:
         missing.append('causal-conv1d')
     if missing:
-        raise ImportError(
-            'TwinkleQwen3_5 linear attention requires flash-linear-attention and causal-conv1d. '
-            f'Missing: {", ".join(missing)}'
-        )
+        raise ImportError('TwinkleQwen3_5 linear attention requires flash-linear-attention and causal-conv1d. '
+                          f'Missing: {", ".join(missing)}')
 
 
 def _maybe_slice_tensor_output(output: Any) -> torch.Tensor:
@@ -66,10 +60,8 @@ def _maybe_slice_tensor_output(output: Any) -> torch.Tensor:
 
 def _sp_is_enabled(sequence_parallel_context: Any | None) -> bool:
     return bool(
-        sequence_parallel_context is not None
-        and getattr(sequence_parallel_context, 'sp_world_size', 1) > 1
-        and getattr(sequence_parallel_context, 'sp_group', None) is not None
-    )
+        sequence_parallel_context is not None and getattr(sequence_parallel_context, 'sp_world_size', 1) > 1
+        and getattr(sequence_parallel_context, 'sp_group', None) is not None)
 
 
 def _get_sp_rank(sequence_parallel_context: Any | None) -> int:
@@ -239,8 +231,7 @@ def _apply_varlen_conv(
     ) -> torch.Tensor:
         if self.causal_conv1d_fn is None:
             raise ImportError(
-                'TwinkleQwen3_5 linear attention requires fla.modules.convolution.causal_conv1d for prefill/train.'
-            )
+                'TwinkleQwen3_5 linear attention requires fla.modules.convolution.causal_conv1d for prefill/train.')
         output = self.causal_conv1d_fn(
             x=mixed_qkv,
             weight=conv_weight,
@@ -261,8 +252,7 @@ def _apply_decode_conv(
         if self.causal_conv1d_update is None:
             raise ImportError(
                 'TwinkleQwen3_5 decode requires a causal_conv1d_update implementation from flash-linear-attention '
-                'or causal-conv1d.'
-            )
+                'or causal-conv1d.')
         mixed_qkv_t = mixed_qkv.transpose(1, 2).contiguous()
         output = self.causal_conv1d_update(
             mixed_qkv_t,
@@ -291,11 +281,8 @@ def forward(
         hidden_states = hf_qwen35.apply_mask_to_padding_states(hidden_states, attention_mask)
         batch_size, seq_len, _ = hidden_states.shape
         use_precomputed_states = (
-            cache_params is not None
-            and cache_params.has_previous_state
-            and seq_len == 1
-            and cache_position is not None
-        )
+            cache_params is not None and cache_params.has_previous_state and seq_len == 1
+            and cache_position is not None)
 
         if cache_params is not None:
             conv_state = cache_params.conv_states[self.layer_idx]
@@ -316,8 +303,7 @@ def forward(
             if self.num_k_heads % sp_world_size != 0 or self.num_v_heads % sp_world_size != 0:
                 raise RuntimeError(
                     'TwinkleQwen3_5 linear attention requires sp_world_size to divide both '
-                    f'linear_num_key_heads ({self.num_k_heads}) and linear_num_value_heads ({self.num_v_heads}).'
-                )
+                    f'linear_num_key_heads ({self.num_k_heads}) and linear_num_value_heads ({self.num_v_heads}).')
             local_num_k_heads = self.num_k_heads // sp_world_size
             local_num_v_heads = self.num_v_heads // sp_world_size
             local_key_dim = local_num_k_heads * self.head_k_dim
@@ -341,7 +327,8 @@ def forward(
                 ),
                 dim=-1,
             )
-            conv_weight = self._get_local_conv1d_weight(_get_sp_rank(sequence_parallel_context), local_key_dim, local_value_dim)
+            conv_weight = self._get_local_conv1d_weight(
+                _get_sp_rank(sequence_parallel_context), local_key_dim, local_value_dim)
         else:
             local_num_k_heads = self.num_k_heads
             local_num_v_heads = self.num_v_heads
@@ -506,8 +493,7 @@ def __init__(self, config: Qwen3_5TextConfig):
         super().__init__(config)
         self.embed_tokens = nn.Embedding(config.vocab_size, config.hidden_size, config.pad_token_id)
         self.layers = nn.ModuleList(
-            [TwinkleQwen3_5DecoderLayer(config, layer_idx) for layer_idx in range(config.num_hidden_layers)]
-        )
+            [TwinkleQwen3_5DecoderLayer(config, layer_idx) for layer_idx in range(config.num_hidden_layers)])
         self.norm = hf_qwen35.Qwen3_5RMSNorm(config.hidden_size, eps=config.rms_norm_eps)
         self.rotary_emb = hf_qwen35.Qwen3_5TextRotaryEmbedding(config=config)
         self.gradient_checkpointing = False
@@ -569,8 +555,7 @@ def forward(
         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
-            )
+                past_seen_tokens, past_seen_tokens + inputs_embeds.shape[1], device=inputs_embeds.device)
 
         if position_ids is None:
             position_ids = cache_position.view(1, 1, -1).expand(3, inputs_embeds.shape[0], -1)

src/twinkle/model/transformers/strategy/sequence_parallel.py

@@ -37,7 +37,8 @@ def get_flattened_cu_seqlens_from_position_ids(position_ids: torch.LongTensor):
         row[row < 0] = 0
         seq_start_indices = torch.where(row == 0)[0]
         if seq_start_indices.numel() == 0 or seq_start_indices[0].item() != 0:
-            seq_start_indices = torch.cat([torch.tensor([0], device=device, dtype=seq_start_indices.dtype), seq_start_indices])
+            seq_start_indices = torch.cat(
+                [torch.tensor([0], device=device, dtype=seq_start_indices.dtype), seq_start_indices])
         seq_end_indices = torch.cat([seq_start_indices[1:], torch.tensor([len(row)], device=device)])
         seq_lengths = (seq_end_indices - seq_start_indices).tolist()
         for seq_length in seq_lengths:
@@ -687,8 +688,7 @@ def prepare(
                 self.causal_mask_func = llm_model._update_causal_mask
         self.attn_implementation = (
             get_config_attr(model.config, '_attn_implementation')
-            or get_config_attr(model.config, '_attn_implementation_internal')
-        )
+            or get_config_attr(model.config, '_attn_implementation_internal'))
 
         if not SequenceParallel._global_inited:
             # these operations are global initializations and patches
@@ -832,8 +832,8 @@ def pad_and_split_inputs(self,
                     cache_position = torch.arange(0, attn_shape, device=inputs.device)
                     # SDPA/eager-style paths still expect a fully materialized causal mask here.
                     if hasattr(self, 'causal_mask_func') and self.causal_mask_func is not None:
-                        attention_mask = self.causal_mask_func(
-                            attention_mask, inputs.to(self.model_dtype), cache_position, None, None)
+                        attention_mask = self.causal_mask_func(attention_mask, inputs.to(self.model_dtype),
+                                                               cache_position, None, None)
         if extra_split_values is not None:
             for (tensor, pad_value, split_dim) in extra_split_values:
                 extra_values.append(

tests/sequence_parallel/test_twinkle_qwen3_5_text_model.py

@@ -1,13 +1,12 @@
 # Copyright (c) ModelScope Contributors. All rights reserved.
 import tempfile
+import torch
 import unittest
 from contextlib import ExitStack
-from types import SimpleNamespace
-from unittest.mock import patch
-
-import torch
 from transformers.models.qwen3_5.configuration_qwen3_5 import Qwen3_5Config, Qwen3_5TextConfig
 from transformers.models.qwen3_5.modeling_qwen3_5 import Qwen3_5ForCausalLM
+from types import SimpleNamespace
+from unittest.mock import patch
 
 from twinkle.model.transformers.models.qwen3_5 import modeling_qwen3_5 as tw_qwen35
 from twinkle.model.transformers.strategy.sequence_parallel import SequenceParallel, SequenceParallelContext
@@ -40,13 +39,9 @@ def _build_text_config(layer_types=None) -> Qwen3_5TextConfig:
 
 
 def _linear_attention_runtime_available() -> bool:
-    return bool(
-        torch.cuda.is_available()
-        and tw_qwen35._FLA_CAUSAL_CONV1D_FN is not None
-        and tw_qwen35._FLA_CHUNK_GATED_DELTA_RULE is not None
-        and tw_qwen35._FLA_FUSED_RECURRENT_GATED_DELTA_RULE is not None
-        and tw_qwen35._HAS_CAUSAL_CONV1D
-    )
+    return bool(torch.cuda.is_available() and tw_qwen35._FLA_CAUSAL_CONV1D_FN is not None
+                and tw_qwen35._FLA_CHUNK_GATED_DELTA_RULE is not None
+                and tw_qwen35._FLA_FUSED_RECURRENT_GATED_DELTA_RULE is not None and tw_qwen35._HAS_CAUSAL_CONV1D)
 
 
 class _ContextReceiver:
@@ -233,13 +228,26 @@ def fake_conv(x, weight, bias, activation, seq_idx=None, backend=None, cu_seqlen
             captured['cu_seqlens'] = cu_seqlens.clone() if cu_seqlens is not None else None
             return x
 
-        def fake_chunk_rule(query, key, value, g, beta, initial_state=None, output_final_state=False,
-                            use_qk_l2norm_in_kernel=False, cu_seqlens=None):
+        def fake_chunk_rule(query,
+                            key,
+                            value,
+                            g,
+                            beta,
+                            initial_state=None,
+                            output_final_state=False,
+                            use_qk_l2norm_in_kernel=False,
+                            cu_seqlens=None):
             del query, key, g, beta, initial_state, output_final_state, use_qk_l2norm_in_kernel
             captured['cu_seqlens'] = cu_seqlens.clone() if cu_seqlens is not None else None
             return value, None
 
-        def fake_recurrent_rule(query, key, value, g, beta, initial_state=None, output_final_state=False,
+        def fake_recurrent_rule(query,
+                                key,
+                                value,
+                                g,
+                                beta,
+                                initial_state=None,
+                                output_final_state=False,
                                 use_qk_l2norm_in_kernel=False):
             del query, key, g, beta, initial_state, output_final_state, use_qk_l2norm_in_kernel
             return value, None
@@ -321,14 +329,27 @@ def fake_conv(x, weight, bias, activation, seq_idx=None, backend=None, cu_seqlen
             captured['cu_seqlens'] = cu_seqlens.clone() if cu_seqlens is not None else None
             return x
 
-        def fake_chunk_rule(query, key, value, g, beta, initial_state=None, output_final_state=False,
-                            use_qk_l2norm_in_kernel=False, cu_seqlens=None):
+        def fake_chunk_rule(query,
+                            key,
+                            value,
+                            g,
+                            beta,
+                            initial_state=None,
+                            output_final_state=False,
+                            use_qk_l2norm_in_kernel=False,
+                            cu_seqlens=None):
             del key, value, g, beta, initial_state, output_final_state, use_qk_l2norm_in_kernel
             captured['query_shape'] = tuple(query.shape)
             captured['cu_seqlens'] = cu_seqlens.clone() if cu_seqlens is not None else None
             return query.new_zeros(query.shape[0], query.shape[1], 4, 4), None
 
-        def fake_recurrent_rule(query, key, value, g, beta, initial_state=None, output_final_state=False,
+        def fake_recurrent_rule(query,
+                                key,
+                                value,
+                                g,
+                                beta,
+                                initial_state=None,
+                                output_final_state=False,
                                 use_qk_l2norm_in_kernel=False):
             del query, key, value, g, beta, initial_state, output_final_state, use_qk_l2norm_in_kernel
             raise AssertionError('recurrent path should not be used')
@@ -368,12 +389,25 @@ def fake_conv(x, weight, bias, activation, seq_idx=None, backend=None, cu_seqlen
             del weight, bias, activation, seq_idx, backend, cu_seqlens
             return x
 
-        def fake_chunk_rule(query, key, value, g, beta, initial_state=None, output_final_state=False,
-                            use_qk_l2norm_in_kernel=False, cu_seqlens=None):
+        def fake_chunk_rule(query,
+                            key,
+                            value,
+                            g,
+                            beta,
+                            initial_state=None,
+                            output_final_state=False,
+                            use_qk_l2norm_in_kernel=False,
+                            cu_seqlens=None):
             del query, key, g, beta, initial_state, output_final_state, use_qk_l2norm_in_kernel, cu_seqlens
             return value, None
 
-        def fake_recurrent_rule(query, key, value, g, beta, initial_state=None, output_final_state=False,
+        def fake_recurrent_rule(query,
+                                key,
+                                value,
+                                g,
+                                beta,
+                                initial_state=None,
+                                output_final_state=False,
                                 use_qk_l2norm_in_kernel=False):
             del query, key, g, beta, initial_state, output_final_state, use_qk_l2norm_in_kernel
             return value, None
@@ -395,8 +429,12 @@ def fake_recurrent_rule(query, key, value, g, beta, initial_state=None, output_f
                     is_packed=False,
                 ))
 
-            def fake_linear_forward(hidden_states, cache_params=None, cache_position=None, attention_mask=None,
-                                    cu_seq_lens_q=None, sequence_parallel_context=None):
+            def fake_linear_forward(hidden_states,
+                                    cache_params=None,
+                                    cache_position=None,
+                                    attention_mask=None,
+                                    cu_seq_lens_q=None,
+                                    sequence_parallel_context=None):
                 del hidden_states, cache_params, cache_position, cu_seq_lens_q, sequence_parallel_context
                 captured['mask'] = attention_mask.clone() if attention_mask is not None else None
                 return torch.zeros(1, 2, config.hidden_size)
@@ -421,8 +459,7 @@ def test_sequence_parallel_drops_dense_attention_mask_for_flash_attention_2(self
         sp.tokenizer = SimpleNamespace(pad_token_id=0)
         sp.model_dtype = torch.bfloat16
         sp.attn_implementation = 'flash_attention_2'
-        sp.causal_mask_func = lambda *args, **kwargs: (_ for _ in ()).throw(
-            AssertionError('should not build 4d mask'))
+        sp.causal_mask_func = lambda *args, **kwargs: (_ for _ in ()).throw(AssertionError('should not build 4d mask'))
 
         input_ids = torch.tensor([[1, 2, 3, 4], [5, 6, 7, 8]], dtype=torch.long)
         position_ids = torch.tensor([[0, 1, 2, 3], [0, 1, 2, 3]], dtype=torch.long)