Add MLA #278
Add MLA #278
Conversation
Config examplemodel:
ddp_bucket_cap_mb: 25
dtype: bfloat16
init_method:
std: 0.025
make_vocab_size_divisible_by: 1
model_config:
...
vocab_size: 50272
# MLA
q_lora_rank: 1536
kv_lora_rank: 512
qk_nope_head_dim: 128
qk_rope_head_dim: 64
v_head_dim: 128 |
| q, [self.qk_nope_head_dim, self.qk_rope_head_dim], dim=-1 | ||
| ) # [seq_len, batch_size, n_local_heads, qk_nope_head_dim], [seq_len, batch_size, n_local_heads, qk_rope_head_dim] | ||
| q_pe = ( | ||
| self.rotary_embedding(q_pe.transpose(0, 1), position_ids=None).transpose(0, 1).contiguous() |
There was a problem hiding this comment.
is the transpose(0,1) needed here?
There was a problem hiding this comment.
yes. otherwise the results would be different
There was a problem hiding this comment.
i meant why not transpose from the beginning of the forward in MLA? this way we avoid doing multiple small transposes
There was a problem hiding this comment.
because transposes are very slow! and you have a lot of them in MLA's forward
| @@ -701,8 +855,9 @@ def __init__( | |||
| layer_idx: int, | |||
| ): | |||
| super().__init__() | |||
| attn_cls = MLA if config.kv_lora_rank is not None else CausalSelfAttention | |||
There was a problem hiding this comment.
I'd rather make it more explicit like use config.use_mla here and assert somewhere that the other configs (e.g. kv_lora_rank) are well defined. This can be done in config.py
There was a problem hiding this comment.
It seems a bit redundant to me since kv_lora_rank = MLA in this case, meaning there's no unexpected behavior
There was a problem hiding this comment.
Sometimes redundancy is fine if it make code cleaner! I still think we should have use_mla somewhere as kv_lora_rank only relates to MLA for now
| q, [self.qk_nope_head_dim, self.qk_rope_head_dim], dim=-1 | ||
| ) # [seq_len, batch_size, n_local_heads, qk_nope_head_dim], [seq_len, batch_size, n_local_heads, qk_rope_head_dim] | ||
| q_pe = ( | ||
| self.rotary_embedding(q_pe.transpose(0, 1), position_ids=None).transpose(0, 1).contiguous() |
There was a problem hiding this comment.
i meant why not transpose from the beginning of the forward in MLA? this way we avoid doing multiple small transposes
| @@ -701,8 +855,9 @@ def __init__( | |||
| layer_idx: int, | |||
| ): | |||
| super().__init__() | |||
| attn_cls = MLA if config.kv_lora_rank is not None else CausalSelfAttention | |||
There was a problem hiding this comment.
Sometimes redundancy is fine if it make code cleaner! I still think we should have use_mla somewhere as kv_lora_rank only relates to MLA for now
| q, [self.qk_nope_head_dim, self.qk_rope_head_dim], dim=-1 | ||
| ) # [seq_len, batch_size, n_local_heads, qk_nope_head_dim], [seq_len, batch_size, n_local_heads, qk_rope_head_dim] | ||
| q_pe = ( | ||
| self.rotary_embedding(q_pe.transpose(0, 1), position_ids=None).transpose(0, 1).contiguous() |
There was a problem hiding this comment.
because transposes are very slow! and you have a lot of them in MLA's forward
| if self.model_config.kv_lora_rank is not None: | ||
| # set num_key_value_heads to None for MLA(as it's same as num_attention_heads in the paper) | ||
| # to avoid unintended errors | ||
| self.model_config.num_key_value_heads = None |
There was a problem hiding this comment.
Please add logger.warning here to warn user
| q_lora_rank: Optional[int] = None | ||
| kv_lora_rank: Optional[int] = None | ||
| qk_nope_head_dim: Optional[int] = None | ||
| qk_rope_head_dim: Optional[int] = None | ||
| v_head_dim: Optional[int] = None |
There was a problem hiding this comment.
we could regroup these in a MLAConfig to make them separate of the rest, or let's just follow transformers' config standards
| ) | ||
|
|
||
| # Initialize linear layers | ||
| self.q_down = nn.Linear(self.dim, self.q_lora_rank, bias=False) # Note: this is duplicated across GPUs |
There was a problem hiding this comment.
Add warning comment please?
Add MLA to Nanotron: