bridge.models.deepseek.attention#

MLA attention spec helpers for the DeepSeek family.

Module Contents#

Classes#

MLASelfAttentionWithoutQueryNorm

MLA self-attention that does not add a query norm when there is no query LoRA.

Functions#

get_deepseek_decoder_block_spec

Build the decoder block spec, omitting the query norm when q_lora_rank is None.

replace_mla_self_attention

Swap MLA self-attention for the query-norm-free variant, in place, on every layer.

API#

class bridge.models.deepseek.attention.MLASelfAttentionWithoutQueryNorm#

Bases: megatron.core.transformer.multi_latent_attention.MLASelfAttention

MLA self-attention that does not add a query norm when there is no query LoRA.

MCore derives Q and KV normalization from a single qk_layernorm flag. DeepSeek needs it enabled for kv_a_layernorm, which every checkpoint ships. When q_lora_rank is None, that same flag also makes MCore fuse a query normalization into linear_q_proj (QKNormConfigResolver._resolve_mla_qk_layernorm), but the HF architecture defines no query-side norm in that case: DeepseekV3Attention builds a bare q_proj.

The result is a trainable parameter with no HF counterpart, which cannot be loaded and is silently dropped on export. This subclass keeps the KV norm and drops the query norm so the converted model matches the source architecture.

Transformer Engine is required for the no-query-LoRA case. MCore builds linear_q_proj from the backend’s fused norm+linear implementation, which only Transformer Engine provides, so the local backend is rejected with an explicit message rather than an internal one.

_resolve_qk_norm_config(submodules)#

Replace the fused query projection with a plain one when there is no query LoRA.

The standalone-q_layernorm case is neutralised before delegating: an MLA spec may set q_layernorm to a real norm whenever qk_layernorm is on, and the parent resolver rejects that outright when there is no query LoRA to consume it (_raise_unused_q_norm). Dropping the query norm is exactly what this class exists to do, so the rejection would fire on a configuration this class already knows how to satisfy.

bridge.models.deepseek.attention.get_deepseek_decoder_block_spec(
config: megatron.core.transformer.transformer_config.TransformerConfig,
use_transformer_engine: bool,
normalization: Optional[str] = None,
qk_l2_norm: Optional[bool] = False,
vp_stage: Optional[int] = None,
pp_rank: Optional[int] = None,
) → megatron.core.transformer.spec_utils.ModuleSpec#

Build the decoder block spec, omitting the query norm when q_lora_rank is None.

The signature mirrors get_gpt_decoder_block_spec exactly, including vp_stage and pp_rank. GPTModelProvider.provide() inspects the callable’s parameters and only forwards vp_stage when it is declared, so dropping it here would leave interleaved pipeline parallelism calling MCore’s layer-offset helper without a virtual stage, which asserts.

Parameters:

  • config – The model provider / transformer config.

  • use_transformer_engine – Whether to build Transformer Engine submodules.

  • normalization – Optional normalization override, forwarded unchanged.

  • qk_l2_norm – Optional QK L2 norm flag, forwarded unchanged.

  • vp_stage – Virtual pipeline stage, forwarded unchanged.

  • pp_rank – Pipeline rank, forwarded unchanged.

Returns:

The decoder block spec, with MLA self-attention replaced by

class:

MLASelfAttentionWithoutQueryNorm when there is no query LoRA.

bridge.models.deepseek.attention.replace_mla_self_attention(
config: megatron.core.transformer.transformer_config.TransformerConfig,
spec: megatron.core.transformer.spec_utils.ModuleSpec,
) → megatron.core.transformer.spec_utils.ModuleSpec#

Swap MLA self-attention for the query-norm-free variant, in place, on every layer.

Shared with the MTP path: a standalone MTP pipeline stage owns no decoder layers, so the provider re-derives a layer spec straight from MCore and never passes through

Func:

get_deepseek_decoder_block_spec. Without this the MTP layer regains the query norm that the decoder layers just dropped.

Accepts either a block spec (.layer_specs) or a single layer spec.