diff --git a/src/mcore_bridge/model/gpts/deepseek_v4.py b/src/mcore_bridge/model/gpts/deepseek_v4.py index ed8a305..f093586 100644 --- a/src/mcore_bridge/model/gpts/deepseek_v4.py +++ b/src/mcore_bridge/model/gpts/deepseek_v4.py @@ -94,6 +94,8 @@ def get_query_key_value_tensors( rotary_pos_emb=None, *, inference_params=None, + boundary_hidden=None, + boundary_rotary_pos_emb=None, ): """ Derives `query`, `key` and `value` tensors from `hidden_states`. @@ -141,7 +143,11 @@ def get_query_key_value_tensors( # QKV up projection and RoPE apply # ========================================= - def qkv_up_proj_and_rope_apply(q_compressed, kv_compressed, rotary_pos_emb): + def qkv_up_proj_and_rope_apply(q_compressed, + kv_compressed, + rotary_pos_emb, + boundary_kv_compressed=None, + boundary_rotary_pos_emb=None): """ Apply the up projection and RoPE to the query and key. When sequence packing enabled, the input tensors adopt a packed shape of [t, ...]; @@ -156,21 +162,18 @@ def qkv_up_proj_and_rope_apply(q_compressed, kv_compressed, rotary_pos_emb): q = q.view(*q.size()[:-1], self.num_attention_heads_per_partition, self.q_head_dim) q = _q_rms_norm(q, self.config.layernorm_epsilon) - kv, _ = self.linear_kv_proj(kv_compressed) - kv = self.kv_layernorm(kv) + boundary_rows = 0 + if boundary_kv_compressed is not None: + boundary_rows = boundary_kv_compressed.shape[0] + kv_projection_input = torch.cat([boundary_kv_compressed, kv_compressed], dim=0) + kv_rotary_pos_emb = torch.cat([boundary_rotary_pos_emb, rotary_pos_emb], dim=0) + else: + kv_projection_input = kv_compressed + kv_rotary_pos_emb = rotary_pos_emb - # [num_tokens, qk_pos_emb_head_dim] -> [num_tokens, 1, qk_pos_emb_head_dim] - q_len = q.size()[0] - if packed_seq_params is None or self.config.context_parallel_size == 1: - # Shorten rotary_pos_emb to the sequence length when inference_params - # is not provided. This makes sure we can run forward directly with - # any sequence length. During training, the sequence length is always - # the full rotary_pos_emb length, except for sequence packing + CP. - # When sequence packing and context parallel are both enabled, the - # position embedding will not split rotary_pos_emb, so it may exceed - # the sequence length on this CP rank, but we need the full rotary_pos_emb - # to cover the full sequence, so we do not shorten it here. - rotary_pos_emb = rotary_pos_emb[0:q_len] + kv, _ = self.linear_kv_proj(kv_projection_input) + kv = self.kv_layernorm(kv) + boundary_kv = None # q_no_pe: [num_tokens, n, qk_head_dim] # q_pos_emb: [num_tokens, n, qk_pos_emb_head_dim] @@ -196,7 +199,7 @@ def qkv_up_proj_and_rope_apply(q_compressed, kv_compressed, rotary_pos_emb): # k_pos_emb:[num_tokens, 1, qk_pos_emb_head_dim] k_pos_emb = apply_rotary_pos_emb( k_pos_emb, - rotary_pos_emb, + kv_rotary_pos_emb, config=self.config, cu_seqlens=cu_seqlens_kv, cp_group=self.pg_collection.cp, @@ -206,24 +209,45 @@ def qkv_up_proj_and_rope_apply(q_compressed, kv_compressed, rotary_pos_emb): # Single head: key = value = [num_tokens, 1, v_head_dim] kv = torch.cat([kv_no_pe, k_pos_emb], dim=-1).unsqueeze(-2) + if boundary_kv_compressed is not None: + boundary_kv = kv[:boundary_rows] + kv = kv[boundary_rows:] key = kv value = kv query = query.contiguous() key = key.contiguous() value = value.contiguous() - - return query, key, value + if boundary_kv is not None: + boundary_kv = boundary_kv.contiguous() + if boundary_kv is None: + return query, key, value + return query, key, value, boundary_kv if self.recompute_up_proj: quantization = self.config.fp8 or self.config.fp4 self.qkv_up_checkpoint = tensor_parallel.CheckpointWithoutOutput(fp8=quantization) - query, key, value = self.qkv_up_checkpoint.checkpoint(qkv_up_proj_and_rope_apply, q_compressed, - kv_compressed, rotary_pos_emb) + if boundary_hidden is None: + query, key, value = self.qkv_up_checkpoint.checkpoint(qkv_up_proj_and_rope_apply, q_compressed, + kv_compressed, rotary_pos_emb) + boundary_kv = None + else: + query, key, value, boundary_kv = self.qkv_up_checkpoint.checkpoint(qkv_up_proj_and_rope_apply, + q_compressed, kv_compressed, + rotary_pos_emb, boundary_hidden, + boundary_rotary_pos_emb) else: - query, key, value = qkv_up_proj_and_rope_apply(q_compressed, kv_compressed, rotary_pos_emb) - - return query, key, value, q_compressed, kv_compressed + if boundary_hidden is None: + query, key, value = qkv_up_proj_and_rope_apply(q_compressed, kv_compressed, rotary_pos_emb) + boundary_kv = None + else: + query, key, value, boundary_kv = qkv_up_proj_and_rope_apply(q_compressed, kv_compressed, rotary_pos_emb, + boundary_hidden, boundary_rotary_pos_emb) + + result = (query, key, value, q_compressed, kv_compressed) + if boundary_kv is not None: + return result + (boundary_kv, ) + return result def forward( self, @@ -251,19 +275,54 @@ def forward( assert (inference_context is None and inference_params is None), 'Inference is not supported for DSv4HybridAttention.' + # Select this microbatch's dynamic CP group. QKV captures it explicitly + # for recompute; the rest of this forward reads it from pg_collection. + # Restore the static group before returning. + cp_group = self.pg_collection.cp + cp_size = cp_group.size() + qkv_format = packed_seq_params.qkv_format if packed_seq_params is not None else None + if cp_size > 1 and qkv_format != 'thd': + raise ValueError("DSv4 Hybrid with CP requires qkv_format='thd'.") + use_thd_cp = cp_size > 1 and qkv_format == 'thd' + if use_thd_cp and packed_seq_params.cp_partition_mode != 'contiguous': + raise ValueError('DSv4 THD CP requires a contiguous CP partition.') + + boundary_hidden = None + boundary_rotary_pos_emb = None + if use_thd_cp: + from megatron.core.transformer.experimental_attention_variant import csa_cp_utils as cp_utils + boundary_hidden = cp_utils.exchange_cp_boundary_hidden( + hidden_states, + self._dsv4_compress_ratio, + self.config.csa_window_size, + self.pg_collection.cp, + ) + boundary_rotary_pos_emb = cp_utils.exchange_cp_boundary_hidden( + rotary_pos_emb, + self._dsv4_compress_ratio, + self.config.csa_window_size, + self.pg_collection.cp, + ) # ===================== # Query, Key, and Value # ===================== # Get the query, key and value tensors based on the type of attention - # self or cross attn. - query, key, value, q_compressed, kv_compressed = self.get_query_key_value_tensors( + qkv = self.get_query_key_value_tensors( hidden_states, key_value_states, position_ids, packed_seq_params, rotary_pos_emb=rotary_pos_emb, inference_context=inference_context, + boundary_hidden=boundary_hidden, + boundary_rotary_pos_emb=boundary_rotary_pos_emb, ) + if use_thd_cp: + query, key, value, q_compressed, kv_compressed, boundary_kv = qkv + else: + query, key, value, q_compressed, kv_compressed = qkv + boundary_kv = None # TODO: Currently, TE can only accept contiguous tensors for MLA query = query.contiguous() @@ -276,6 +335,10 @@ def forward( # Need corresponding TE change core_attn_manager = off_interface(self.offload_core_attention and self.training, query, 'core_attn') with core_attn_manager as query: + core_attn_kwargs = {} + if boundary_hidden is not None: + core_attn_kwargs['boundary_hidden'] = boundary_hidden + core_attn_kwargs['boundary_kv'] = boundary_kv core_attn_out = self.core_attention( query, key, @@ -284,8 +347,12 @@ def forward( packed_seq_params=packed_seq_params, x=hidden_states, qr=q_compressed, + **core_attn_kwargs, ) - core_attn_out = core_attn_manager.group_offload(core_attn_out, forced_released_tensors=[query, key, value]) + forced_released_tensors = [query, key, value] + if boundary_kv is not None: + forced_released_tensors.append(boundary_kv) + core_attn_out = core_attn_manager.group_offload(core_attn_out, forced_released_tensors=forced_released_tensors) if packed_seq_params is not None and packed_seq_params.qkv_format == 'thd': # reshape to same output shape as unpacked case @@ -313,8 +380,12 @@ def forward( cu_seqlens_kv = None content_part, rot_part = torch.split(core_attn_out, [core_attn_out.size(-1) - pos_dim, pos_dim], dim=-1) - rot_part = apply_rotary_pos_emb( - rot_part, + if packed_seq: + rot_part_in = rot_part.squeeze(1) + else: + rot_part_in = rot_part + rot_part_out = apply_rotary_pos_emb( + rot_part_in, rotary_pos_emb, self.config, cu_seqlens=cu_seqlens_kv, @@ -323,6 +394,10 @@ def forward( inverse=True, mla_output_remove_interleaving=True, ) + if packed_seq: + rot_part = rot_part_out.unsqueeze(1) + else: + rot_part = rot_part_out core_attn_out = torch.cat([content_part, rot_part], dim=-1) core_attn_out = core_attn_out.view(seq_len, core_attn_out.size(1), -1) diff --git a/src/mcore_bridge/patcher.py b/src/mcore_bridge/patcher.py index 37390ac..4909e3a 100644 --- a/src/mcore_bridge/patcher.py +++ b/src/mcore_bridge/patcher.py @@ -203,7 +203,9 @@ def _apply_rotary_pos_emb_thd(t: torch.Tensor, cu_seqlens: torch.Tensor, freqs: cp_size = mpu.get_context_parallel_world_size() cu_seqlens_for_batched = cu_seqlens // cp_size use_batched_rope = (freqs.dim() >= 1 and freqs.shape[0] == cu_seqlens_for_batched[-1]).item() - if not use_batched_rope: + # The determination of mla_output_remove_interleaving: a quick solution for identifying deepseek_v4 + # (TODO: refactor) + if not use_batched_rope and not kwargs.get('mla_output_remove_interleaving', False): logger.warning_once('Using non-batched RoPE, which may affect performance.') return _origin_apply_rotary_pos_emb_thd(t, cu_seqlens, freqs, *args, **kwargs) diff --git a/src/mcore_bridge/utils/megatron_utils.py b/src/mcore_bridge/utils/megatron_utils.py index 0de33c1..bdaf50b 100644 --- a/src/mcore_bridge/utils/megatron_utils.py +++ b/src/mcore_bridge/utils/megatron_utils.py @@ -45,12 +45,26 @@ def unwrap_model(models, module_instances=None): return unwrapped_model -def split_cp_inputs(inputs: torch.Tensor, cu_seqlens: Optional[torch.Tensor], dim: int): +def split_cp_inputs(inputs: torch.Tensor, + cu_seqlens: Optional[torch.Tensor], + dim: int, + cp_partition_mode: str = 'zigzag'): if dim < 0: dim = (dim + inputs.ndim) % inputs.ndim - new_inputs = [] cp_size = mpu.get_context_parallel_world_size() cp_rank = mpu.get_context_parallel_rank() + if cp_partition_mode == 'contiguous': + # Rank r owns the contiguous block [r * local_rows, (r + 1) * local_rows) of the + # flattened packed sequence. This must match the DSv4 THD CP forward, which assumes + # each rank holds a single contiguous slice (global_start = cp_rank * local_rows). + total_rows = inputs.shape[dim] + assert total_rows % cp_size == 0, ( + f'Contiguous CP slicing requires dim size={total_rows} to be divisible by cp_size={cp_size}.') + local_rows = total_rows // cp_size + slices = [slice(None)] * inputs.ndim + slices[dim] = slice(cp_rank * local_rows, (cp_rank + 1) * local_rows) + return inputs[tuple(slices)].contiguous() + new_inputs = [] for i in range(1 if cu_seqlens is None else (cu_seqlens.shape[0] - 1)): if cu_seqlens is None: val = inputs