diff --git a/data/v1_multimodal_demo.jsonl b/data/v1_multimodal_demo.jsonl index 2094834f6..d6cd0b1bb 100644 --- a/data/v1_multimodal_demo.jsonl +++ b/data/v1_multimodal_demo.jsonl @@ -22,4 +22,3 @@ {"messages": [{"role": "user", "content": [{"type": "image_url", "value": "data/mllm_demo_data/1.jpg"}, {"type": "text", "value": "他们是谁？"}]}, {"role": "assistant", "content": [{"type": "text", "value": "他们是拜仁慕尼黑的凯恩和格雷茨卡。"}]}, {"role": "user", "content": [{"type": "text", "value": "他们在做什么？"}, {"type": "image_url", "value": "data/mllm_demo_data/1.jpg"}]}, {"role": "assistant", "content": [{"type": "text", "value": "他们在足球场上庆祝。"}]}]} {"messages": [{"role": "user", "content": [{"type": "image_url", "value": "data/mllm_demo_data/2.jpg"}, {"type": "text", "value": "他是谁？"}]}, {"role": "assistant", "content": [{"type": "text", "value": "他是来自拜仁慕尼黑的托马斯·穆勒。"}]}, {"role": "user", "content": [{"type": "text", "value": "他为什么在地上？"}]}, {"role": "assistant", "content": [{"type": "text", "value": "因为他正在双膝跪地滑行庆祝。"}]}]} {"messages": [{"role": "user", "content": [{"type": "image_url", "value": "data/mllm_demo_data/3.jpg"}, {"type": "text", "value": "请描述这张图片"}]}, {"role": "assistant", "content": [{"type": "text", "value": "中国宇航员桂海潮正在讲话。"}]}, {"role": "user", "content": [{"type": "text", "value": "他取得过哪些成就？"}]}, {"role": "assistant", "content": [{"type": "text", "value": "他于2022年6月被任命为神舟十六号任务的有效载荷专家，从而成为2023年5月30日进入太空的首位平民宇航员。他负责在轨操作空间科学实验有效载荷。"}]}]} - diff --git a/data/v1_multimodal_demo.yaml b/data/v1_multimodal_demo.yaml index 74189bc1f..7ecb0a454 100644 --- a/data/v1_multimodal_demo.yaml +++ b/data/v1_multimodal_demo.yaml @@ -1,4 +1,3 @@ multimodal_demo: path: data/v1_multimodal_demo.jsonl source: local - diff --git a/examples/v1/train_full/train_multimodal.yaml b/examples/v1/train_full/train_multimodal.yaml index 07f8a980c..084e041ed 100644 --- a/examples/v1/train_full/train_multimodal.yaml +++ b/examples/v1/train_full/train_multimodal.yaml @@ -24,4 +24,3 @@ max_steps: 5 ### sample sample_backend: hf max_new_tokens: 128 - diff --git a/src/llamafactory/v1/core/base_trainer.py b/src/llamafactory/v1/core/base_trainer.py index 8b4cfcbe4..5588f67e2 100644 --- a/src/llamafactory/v1/core/base_trainer.py +++ b/src/llamafactory/v1/core/base_trainer.py @@ -147,11 +147,6 @@ class BaseTrainer: self.state.epoch = self._resume_epoch if self.args.cp_size > 1: - # qwen3.5 is not supported because of the different attention implementation, which will be supported in the future. - if model.config.model_type == "qwen3_5": - raise RuntimeError( - "Sequence parallel is not supported for qwen3.5 model due to its different attention implementation, which will be supported in the future." - ) from ..plugins.model_plugins.parallelization.sequence_parallel import SequenceParallelModelPlugin if model.config._attn_implementation != "flash_attention_2": diff --git a/src/llamafactory/v1/core/rendering/rendering.py b/src/llamafactory/v1/core/rendering/rendering.py index 94aa128d1..87f7d6a76 100644 --- a/src/llamafactory/v1/core/rendering/rendering.py +++ b/src/llamafactory/v1/core/rendering/rendering.py @@ -23,7 +23,6 @@ Note: ``position_ids`` are assigned by ``process_samples`` (1-based); multimodal ids are expected to be recomputed by the model/trainer. """ - import json import numpy as np diff --git a/src/llamafactory/v1/plugins/data_plugins/converter.py b/src/llamafactory/v1/plugins/data_plugins/converter.py index 556f83ca1..0882e0339 100644 --- a/src/llamafactory/v1/plugins/data_plugins/converter.py +++ b/src/llamafactory/v1/plugins/data_plugins/converter.py @@ -295,4 +295,3 @@ def pair_converter(raw_sample: PairSample) -> DPOSample: logger.warning_rank0(f"Invalid tools format: {str(tools)}") return sample - diff --git a/src/llamafactory/v1/plugins/model_plugins/parallelization/gdn_attention.py b/src/llamafactory/v1/plugins/model_plugins/parallelization/gdn_attention.py new file mode 100644 index 000000000..9f28249f5 --- /dev/null +++ b/src/llamafactory/v1/plugins/model_plugins/parallelization/gdn_attention.py @@ -0,0 +1,250 @@ +# Copyright 2025 the LlamaFactory team. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + + +import torch +import torch.distributed as dist +import torch.nn.functional as F + +from .seq_comm import SeqAllToAll4D +from .ulysses import ( + get_ulysses_sequence_parallel_group, + get_ulysses_sequence_parallel_world_size, +) + + +def is_gdn_layer(layer) -> bool: + """Return True if the module is a GDN (linear attention) layer or a DecoderLayer containing one.""" + if hasattr(layer, "layer_type") and layer.layer_type == "linear_attention": + return True + if hasattr(layer, "block_type") and layer.block_type == "linear_attention": + return True + return False + + +def _get_gdn_module(module): + """Return the actual GDN module from either a GDN layer or a DecoderLayer.""" + if hasattr(module, "in_proj_qkv"): + return module + if hasattr(module, "linear_attn"): + return module.linear_attn + raise AttributeError( + f"Cannot find GDN module on {type(module).__name__}. " + f"Expected either a GDN layer with in_proj_qkv or a DecoderLayer with linear_attn." + ) + + +def get_parameter_local_cp(param, dim, cp_group, split_sections=None): + """Slice a parameter for the current CP rank. + + If split_sections is given, first split along dim into sub-groups, + slice each sub-group independently for CP, then concatenate back. + This ensures each CP rank gets a proportional slice of each sub-group. + """ + cp_size = dist.get_world_size(group=cp_group) + if cp_size == 1: + return param + cp_rank = dist.get_rank(group=cp_group) + + if split_sections is not None: + inputs = torch.split(param, split_sections, dim=dim) + outputs = [] + for p in inputs: + p = get_parameter_local_cp(p, dim, cp_group) + outputs.append(p) + return torch.cat(outputs, dim=dim) + + slices = [slice(None)] * param.dim() + dim_size = param.size(dim=dim) + slices[dim] = slice(cp_rank * dim_size // cp_size, (cp_rank + 1) * dim_size // cp_size) + return param[slices] + + +def gdn_forward_with_cp(self, hidden_states, attention_mask=None, **kwargs): + """GDN forward with Context Parallel support. + + Uses SeqAllToAll4D (same as UlyssesAttention) for all all_to_all operations. + Each component (Q/K/V/z/b/a) is independently reshaped to 4D and all_to_all'd + with scatter heads / gather seq, avoiding the bug where uniform hidden-split on + merged qkv gives rank-0 [Q+K] and rank-1 [V]. + + Falls back to self.original_forward when cp_size <= 1. + """ + cp_size = get_ulysses_sequence_parallel_world_size() + if cp_size <= 1: + return self.original_forward(hidden_states, attention_mask=attention_mask, **kwargs) + + cp_group = get_ulysses_sequence_parallel_group() + + if attention_mask is not None: + try: + from transformers.models.qwen3_5.modeling_qwen3_5 import apply_mask_to_padding_states + + hidden_states = apply_mask_to_padding_states(hidden_states, attention_mask) + except ImportError: + pass + + batch_size, seq_len, _ = hidden_states.shape + full_seq_len = seq_len * cp_size + + # Extract position_ids and derive cu_seqlens for pack support + position_ids = kwargs.get("position_ids", None) + cu_seqlens = None + if position_ids is not None and batch_size == 1: + global_position_ids = [torch.empty_like(position_ids) for _ in range(cp_size)] + dist.all_gather(global_position_ids, position_ids, group=cp_group) + global_position_ids = torch.cat(global_position_ids, dim=-1).contiguous() + try: + from transformers.modeling_flash_attention_utils import prepare_fa_kwargs_from_position_ids + cu_seqlens = prepare_fa_kwargs_from_position_ids(global_position_ids)[0][0] + except ImportError: + cu_seqlens = None + + # Input projections in CP layout: [B, seq/cp, hidden] + qkv = self.in_proj_qkv(hidden_states) # [B, seq/cp, key_dim*2 + value_dim] + z = self.in_proj_z(hidden_states) # [B, seq/cp, value_dim] + b = self.in_proj_b(hidden_states) # [B, seq/cp, num_v_heads] + a = self.in_proj_a(hidden_states) # [B, seq/cp, num_v_heads] + + # Split qkv into Q, K, V before all_to_all so each sub-group gets + # proportional head distribution across ranks. + q_proj, k_proj, v_proj = torch.split(qkv, [self.key_dim, self.key_dim, self.value_dim], dim=-1) + + # CP->HP all_to_all for each component: scatter heads (dim=2), gather seq (dim=1) + # [B, S/cp, heads, head_dim] -> [B, S, heads/cp, head_dim] + q_proj = q_proj.reshape(batch_size, seq_len, self.num_k_heads, self.head_k_dim) + q_proj = SeqAllToAll4D.apply(cp_group, q_proj, 2, 1) + + k_proj = k_proj.reshape(batch_size, seq_len, self.num_k_heads, self.head_k_dim) + k_proj = SeqAllToAll4D.apply(cp_group, k_proj, 2, 1) + + v_proj = v_proj.reshape(batch_size, seq_len, self.num_v_heads, self.head_v_dim) + v_proj = SeqAllToAll4D.apply(cp_group, v_proj, 2, 1) + + z = z.reshape(batch_size, seq_len, self.num_v_heads, self.head_v_dim) + z = SeqAllToAll4D.apply(cp_group, z, 2, 1) + + b = b.reshape(batch_size, seq_len, self.num_v_heads, 1) + b = SeqAllToAll4D.apply(cp_group, b, 2, 1) + + a = a.reshape(batch_size, seq_len, self.num_v_heads, 1) + a = SeqAllToAll4D.apply(cp_group, a, 2, 1) + + # Merge Q/K/V for conv1d (conv1d requires merged qkv) + q_flat = q_proj.reshape(batch_size, full_seq_len, self.key_dim // cp_size) + k_flat = k_proj.reshape(batch_size, full_seq_len, self.key_dim // cp_size) + v_flat = v_proj.reshape(batch_size, full_seq_len, self.value_dim // cp_size) + qkv = torch.cat([q_flat, k_flat, v_flat], dim=-1) # [B, S, (key_dim*2+value_dim)/cp] + + # Conv1d in HP layout with CP-aware weight slicing + mixed_qkv = qkv.transpose(1, 2).contiguous() # [B, conv_dim/cp, S] + conv1d_weight = get_parameter_local_cp( + self.conv1d.weight, + dim=0, + cp_group=cp_group, + split_sections=[self.key_dim, self.key_dim, self.value_dim], + ) + conv1d_bias = None + if self.conv1d.bias is not None: + conv1d_bias = get_parameter_local_cp( + self.conv1d.bias, + dim=0, + cp_group=cp_group, + split_sections=[self.key_dim, self.key_dim, self.value_dim], + ) + + if self.causal_conv1d_fn is not None: + mixed_qkv = self.causal_conv1d_fn( + x=mixed_qkv, + weight=conv1d_weight.squeeze(1), + bias=conv1d_bias, + activation=self.activation, + seq_idx=None, + **({"cu_seqlens": cu_seqlens} if cu_seqlens is not None else {}), + ) + elif cu_seqlens is not None: + raise RuntimeError( + "cu_seqlens requires causal_conv1d_fn (FLA) but it is not available. " + "Please install flash-linear-attention for pack support." + ) + else: + conv_out = F.conv1d( + input=mixed_qkv, + weight=conv1d_weight, + bias=conv1d_bias, + stride=self.conv1d.stride, + padding=self.conv1d.padding, + dilation=self.conv1d.dilation, + groups=self.conv_dim // cp_size, + ) + mixed_qkv = self.act(conv_out[..., :full_seq_len]) + mixed_qkv = mixed_qkv.transpose(1, 2).contiguous() # [B, S, conv_dim/cp] + + query, key, value = torch.split( + mixed_qkv, + [self.key_dim // cp_size, self.key_dim // cp_size, self.value_dim // cp_size], + dim=-1, + ) + query = query.reshape(batch_size, full_seq_len, -1, self.head_k_dim) + key = key.reshape(batch_size, full_seq_len, -1, self.head_k_dim) + value = value.reshape(batch_size, full_seq_len, -1, self.head_v_dim) + + if self.num_v_heads // self.num_k_heads > 1: + repeat_factor = self.num_v_heads // self.num_k_heads + query = query.repeat_interleave(repeat_factor, dim=2) + key = key.repeat_interleave(repeat_factor, dim=2) + + gate = z # [B, S, num_v_heads/cp, head_v_dim] + beta = b.squeeze(-1) # [B, S, num_v_heads/cp] + alpha = a.squeeze(-1) # [B, S, num_v_heads/cp] + + query = query.contiguous() + key = key.contiguous() + value = value.contiguous() + gate = gate.contiguous() + beta = beta.contiguous() + alpha = alpha.contiguous() + + A_log_local = get_parameter_local_cp(self.A_log, dim=0, cp_group=cp_group) + dt_bias_local = get_parameter_local_cp(self.dt_bias, dim=0, cp_group=cp_group) + g = -A_log_local.float().exp() * F.softplus(alpha.float() + dt_bias_local) + beta_final = beta.sigmoid() + + # Gated delta rule in HP layout (needs full sequence) + core_attn_out, _ = self.chunk_gated_delta_rule( + query, + key, + value, + g=g, + beta=beta_final, + initial_state=None, + output_final_state=False, + use_qk_l2norm_in_kernel=True, + **({"cu_seqlens": cu_seqlens} if cu_seqlens is not None else {}), + ) + + z_shape_og = gate.shape + core_attn_out = core_attn_out.reshape(-1, core_attn_out.shape[-1]) + z_flat = gate.reshape(-1, gate.shape[-1]) + core_attn_out = self.norm(core_attn_out, z_flat) + core_attn_out = core_attn_out.reshape(z_shape_og) # [B, S, num_v_heads/cp, head_v_dim] + + # HP->CP all_to_all: scatter seq (dim=1), gather heads (dim=2) + # [B, S, num_v_heads/cp, head_v_dim] -> [B, S/cp, num_v_heads, head_v_dim] + norm_out = SeqAllToAll4D.apply(cp_group, core_attn_out, 1, 2) + norm_out = norm_out.reshape(batch_size, seq_len, -1) + + # Output projection in CP layout + output = self.out_proj(norm_out) + return output diff --git a/src/llamafactory/v1/plugins/model_plugins/parallelization/sequence_parallel.py b/src/llamafactory/v1/plugins/model_plugins/parallelization/sequence_parallel.py index 47a6a1d99..517d9a7a4 100644 --- a/src/llamafactory/v1/plugins/model_plugins/parallelization/sequence_parallel.py +++ b/src/llamafactory/v1/plugins/model_plugins/parallelization/sequence_parallel.py @@ -24,6 +24,7 @@ from ....accelerator.interface import Dim, DistributedInterface from ....utils import logging from ....utils.plugin import BasePlugin from ....utils.types import ModelOutput +from .gdn_attention import _get_gdn_module, gdn_forward_with_cp, is_gdn_layer from .ulysses import ( UlyssesAttention, get_ulysses_sequence_parallel_group, @@ -127,6 +128,20 @@ def apply_sequence_parallel(model, cp_size: int): except (AttributeError, TypeError): continue + # Register GDN forward for CP support + if cp_size > 1: + replaced_modules = set() + for name, module in model.named_modules(): + if is_gdn_layer(module): + gdn_module = _get_gdn_module(module) + if id(gdn_module) in replaced_modules: + continue + replaced_modules.add(id(gdn_module)) + gdn_module.original_forward = gdn_module.forward + gdn_module.forward = gdn_forward_with_cp.__get__(gdn_module, type(gdn_module)) + gdn_name = name if gdn_module is module else f"{name}.linear_attn" + logger.info_rank0(f"Replaced GDN forward in {gdn_name} with gdn_forward_with_cp for context parallel.") + def padding_and_split_data(data, device_mesh=None): if device_mesh is not None: diff --git a/tests_v1/core/utils/test_batching.py b/tests_v1/core/utils/test_batching.py index 2767e175e..9da38ef8a 100644 --- a/tests_v1/core/utils/test_batching.py +++ b/tests_v1/core/utils/test_batching.py @@ -566,12 +566,11 @@ def test_drop_unsupervised_samples(): gen = SimpleNamespace(cutoff_len=4, _warned_truncation=False) samples = [ - _s([0.0, 0.0, 1.0, 1.0]), # fits cutoff (len 4), supervised -> kept - _s([0.0, 0.0, 0.0, 0.0, 1.0, 1.0]), # len 6 > 4, supervision only beyond cutoff -> dropped - _s([1.0, 1.0]), # short, fully supervised -> kept - _s([0.0, 0.0, 1.0, 1.0, 1.0, 1.0]), # len 6 > 4 but supervision within cutoff -> kept + _s([0.0, 0.0, 1.0, 1.0]), # fits cutoff (len 4), supervised -> kept + _s([0.0, 0.0, 0.0, 0.0, 1.0, 1.0]), # len 6 > 4, supervision only beyond cutoff -> dropped + _s([1.0, 1.0]), # short, fully supervised -> kept + _s([0.0, 0.0, 1.0, 1.0, 1.0, 1.0]), # len 6 > 4 but supervision within cutoff -> kept ] kept = BatchGenerator._drop_unsupervised(gen, samples) assert kept == [samples[0], samples[2], samples[3]] assert gen._warned_truncation is True - diff --git a/tests_v1/plugins/data_plugins/test_converter.py b/tests_v1/plugins/data_plugins/test_converter.py index 20ddee955..7d836477b 100644 --- a/tests_v1/plugins/data_plugins/test_converter.py +++ b/tests_v1/plugins/data_plugins/test_converter.py @@ -259,4 +259,3 @@ def test_pair_converter(num_samples: int): ], } assert data_engine[index] == {"_dataset_name": "tiny_dataset", **expected_data} -