diff --git a/src/mcore_bridge/config/parser.py b/src/mcore_bridge/config/parser.py index badaac3..ece484c 100644 --- a/src/mcore_bridge/config/parser.py +++ b/src/mcore_bridge/config/parser.py @@ -213,6 +213,8 @@ def hf_to_mcore_config(hf_config: PretrainedConfig) -> Dict[str, Any]: res.setdefault('linear_attention_freq', 4) elif llm_model_type == 'minimax_m2': res['add_qkv_bias'] = False + elif llm_model_type == 'olmoe': + res['qk_layernorm'] = True elif hf_model_type == 'llama4': qk_layernorm = res.pop('qk_layernorm', False) if qk_layernorm: diff --git a/src/mcore_bridge/model/gpts/olmoe.py b/src/mcore_bridge/model/gpts/olmoe.py index 63c1957..dd30ef3 100644 --- a/src/mcore_bridge/model/gpts/olmoe.py +++ b/src/mcore_bridge/model/gpts/olmoe.py @@ -1,222 +1,95 @@ -import torch -import torch.distributed as dist -from copy import deepcopy -from megatron.core.extensions.transformer_engine import SplitAlongDim, TENorm -from megatron.core.models.gpt.gpt_layer_specs import get_gpt_layer_with_transformer_engine_spec -from megatron.core.transformer.attention import SelfAttention as SelfAttentionBase -from megatron.core.transformer.attention import SelfAttentionSubmodules +# Copyright (c) ModelScope Contributors. All rights reserved. +from megatron.core import mpu +from megatron.core.tensor_parallel.mappings import (gather_from_tensor_model_parallel_region, + scatter_to_tensor_model_parallel_region) +from megatron.core.transformer.attention import SelfAttention, SelfAttentionSubmodules +from megatron.core.transformer.identity_op import IdentityOp from megatron.core.transformer.spec_utils import build_module -from megatron.core.transformer.transformer_block import TransformerBlockSubmodules, get_num_layers_to_build -from megatron.core.transformer.transformer_layer import get_transformer_layer_offset from typing import Optional from mcore_bridge.bridge import GPTBridge from mcore_bridge.config import ModelConfig -from mcore_bridge.tuners import LoraParallelLinear from ..constant import ModelType from ..register import ModelLoader, ModelMeta, register_model -class OLMoESelfAttention(SelfAttentionBase): - - def __init__(self, config: ModelConfig, submodules: SelfAttentionSubmodules, *args, **kwargs): - super().__init__(config, submodules, *args, **kwargs) - self.q_layernorm = build_module( +class OLMoESelfAttention(SelfAttention): + + def __init__( + self, + config: ModelConfig, + submodules: SelfAttentionSubmodules, + *args, + **kwargs, + ): + q_layernorm = submodules.q_layernorm + k_layernorm = submodules.k_layernorm + submodules.q_layernorm = IdentityOp + submodules.k_layernorm = IdentityOp + try: + super().__init__(config, submodules, *args, **kwargs) + finally: + submodules.q_layernorm = q_layernorm + submodules.k_layernorm = k_layernorm + # OLMoE applies q/k RMSNorm over the full Q/K channels (not per-head). + self.q_norm = build_module( submodules.q_layernorm, - hidden_size=self.hidden_size_per_attention_head * self.num_attention_heads_per_partition, + hidden_size=self.hidden_size_per_attention_head * config.num_attention_heads, config=self.config, eps=self.config.layernorm_epsilon, ) - self.k_layernorm = build_module( + self.k_norm = build_module( submodules.k_layernorm, - hidden_size=self.hidden_size_per_attention_head * self.num_query_groups_per_partition, + hidden_size=self.hidden_size_per_attention_head * config.num_query_groups, config=self.config, eps=self.config.layernorm_epsilon, ) - - def get_query_key_value_tensors(self, hidden_states, key_value_states=None, *args, **kwargs): - """ - Derives `query`, `key` and `value` tensors from `hidden_states`. - """ - # Attention heads [sq, b, h] --> [sq, b, ng * (np/ng + 2) * hn)] - mixed_qkv, _ = self.linear_qkv(hidden_states) - - # [sq, b, ng * (np/ng + 2) * hn] -> [sq, b, np * hn], [sq, b, ng * hn], [sq, b, ng * hn] - split_arg_list = [ - self.hidden_size_per_attention_head * self.num_attention_heads_per_partition, - self.hidden_size_per_attention_head * self.num_query_groups_per_partition, - self.hidden_size_per_attention_head * self.num_query_groups_per_partition - ] - - if SplitAlongDim is not None: - (query, key, value) = SplitAlongDim(mixed_qkv, 2, split_arg_list) - else: - (query, key, value) = torch.split(mixed_qkv, split_arg_list, dim=2) - if self.q_layernorm is not None: - query = self.q_layernorm(query) - - if self.k_layernorm is not None: - key = self.k_layernorm(key) - query = query.reshape(query.size(0), query.size(1), -1, self.hidden_size_per_attention_head) - key = key.reshape(key.size(0), key.size(1), -1, self.hidden_size_per_attention_head) - value = value.reshape(value.size(0), value.size(1), -1, self.hidden_size_per_attention_head) - - if self.config.test_mode: - self.run_realtime_tests() - + # These norms operate on the gathered (TP-replicated) tensor, so they must + # NOT participate in the cross-TP SP/QK-layernorm grad AllReduce, + # otherwise the gradients will be amplified by tp_size. + for norm in (self.q_norm, self.k_norm): + for p in norm.parameters(): + p.sequence_parallel = False + p.average_gradients_across_tp_domain = True + + def get_query_key_value_tensors(self, *_args, **kwargs): + enable_tp = mpu.get_tensor_model_parallel_world_size() > 1 + query, key, value = super().get_query_key_value_tensors(*_args, **kwargs) + query = query.reshape(*query.shape[:-2], -1) + key = key.reshape(*key.shape[:-2], -1) + if enable_tp: + query = gather_from_tensor_model_parallel_region(query) + key = gather_from_tensor_model_parallel_region(key) + query = self.q_norm(query) + key = self.k_norm(key) + if enable_tp: + query = scatter_to_tensor_model_parallel_region(query) + key = scatter_to_tensor_model_parallel_region(key) + query = query.view(*query.shape[:2], -1, self.hidden_size_per_attention_head) + key = key.view(*key.shape[:2], -1, self.hidden_size_per_attention_head) return query, key, value -def get_olmoe_decoder_block_spec( - config: ModelConfig, - vp_stage: Optional[int] = None, -) -> TransformerBlockSubmodules: - """GPT block spec.""" - layer_norm_impl = TENorm - moe_layer_spec = get_gpt_layer_with_transformer_engine_spec( - num_experts=config.num_moe_experts, - moe_grouped_gemm=config.moe_grouped_gemm, - qk_layernorm=True, - multi_latent_attention=False, - use_kitchen=config.use_kitchen, - ) - layer_specs = [] - for _ in range(config.num_layers): - layer_spec = deepcopy(moe_layer_spec) - layer_spec.submodules.self_attention.module = OLMoESelfAttention - layer_specs.append(layer_spec) - - num_layers_to_build = get_num_layers_to_build(config, vp_stage=vp_stage) - - if config.pipeline_model_parallel_layout is not None: - from megatron.core.transformer.enums import LayerType - local_layer_specs = [ - layer_specs[layer_id] for layer_id in config.pipeline_model_parallel_layout.get_layer_id_list( - layer_type=LayerType.decoder, vp_stage=vp_stage) - ] - else: - offset = get_transformer_layer_offset(config, vp_stage=vp_stage) - local_layer_specs = layer_specs[offset:offset + num_layers_to_build] - - # Block spec. - block_spec = TransformerBlockSubmodules(layer_specs=local_layer_specs, layer_norm=layer_norm_impl) - - return block_spec - - class OLMoEBridge(GPTBridge): - def _set_attn_state(self, mg_attn, hf_state_dict, hf_prefix: str, layer_idx: int, to_mcore: bool): - if to_mcore: - hf_state_dict = self._remove_prefix(hf_state_dict, hf_prefix) - else: - hf_state_dict = {} - config = self.config - if to_mcore: - if isinstance(mg_attn.linear_qkv, LoraParallelLinear): - lora_A = hf_state_dict['q_proj.lora_A.weight'].load() - assert (lora_A == hf_state_dict['k_proj.lora_A.weight'].load()).all() and ( - lora_A == hf_state_dict['v_proj.lora_A.weight'].load() - ).all(), 'Need to ensure QKV\'s lora_A are consistent' - lora_B = torch.cat([ - hf_state_dict['q_proj.lora_B.weight'].load(), - hf_state_dict['k_proj.lora_B.weight'].load(), - hf_state_dict['v_proj.lora_B.weight'].load(), - ], - dim=0) - self._set_weight(mg_attn.linear_qkv.lora_A[self._adapter_name].weight, lora_A, - 'linear_qkv.lora_A.weight') - self._set_weight(mg_attn.linear_qkv.lora_B[self._adapter_name].weight, lora_B, - 'linear_qkv.lora_B.weight') - elif not self._peft_format: - linear_qkv_weight = torch.cat([ - hf_state_dict['q_proj.weight'].load(), - hf_state_dict['k_proj.weight'].load(), - hf_state_dict['v_proj.weight'].load(), - ], - dim=0) - qkv_scale_inv = None - if 'q_proj.weight_scale_inv' in hf_state_dict: - qkv_scale_inv = torch.cat([ - hf_state_dict['q_proj.weight_scale_inv'].load(), - hf_state_dict['k_proj.weight_scale_inv'].load(), - hf_state_dict['v_proj.weight_scale_inv'].load(), - ], - dim=0) - self._set_weight( - mg_attn.linear_qkv.weight, linear_qkv_weight, 'linear_qkv.weight', hf_scale_inv=qkv_scale_inv) - else: - q_dim = self.config.kv_channels * self.config.num_attention_heads // self.config.num_query_groups - kv_dim = self.config.kv_channels - q_block = q_dim // self.fp8_block_size - kv_block = kv_dim // self.fp8_block_size - is_lora = False if mg_attn is None else isinstance(mg_attn.linear_qkv, - LoraParallelLinear) and self._peft_format - is_lora = torch.tensor([is_lora], dtype=torch.bool, device='cuda') - if self.pp_size > 1: - dist.all_reduce(is_lora, group=self.pp_group) - if is_lora: - lora_A, _ = self._get_weight( - None if mg_attn is None else mg_attn.linear_qkv.lora_A[self._adapter_name].weight.data, - f'linear_qkv.lora_A.{self._adapter_name}.weight') - lora_B, _ = self._get_weight( - None if mg_attn is None else mg_attn.linear_qkv.lora_B[self._adapter_name].weight.data, - f'linear_qkv.lora_B.{self._adapter_name}.weight') - if lora_A is not None: - self._peft_target_modules.update({'q_proj', 'k_proj', 'v_proj'}) - for key in ['q_proj', 'k_proj', 'v_proj']: - hf_state_dict[f'{key}.lora_A.weight'] = lora_A.clone() - hf_state_dict['q_proj.lora_B.weight'] = lora_B[:q_dim, :].clone() - hf_state_dict['k_proj.lora_B.weight'] = lora_B[q_dim:-kv_dim, :].clone() - hf_state_dict['v_proj.lora_B.weight'] = lora_B[-kv_dim:, :].clone() - elif not self._peft_format: - mg_attn_weight, scale_inv = self._get_weight( - None if mg_attn is None else mg_attn.linear_qkv.weight.data, 'linear_qkv.weight') - if mg_attn_weight is not None: - hf_state_dict['q_proj.weight'] = mg_attn_weight[:q_dim, :].clone() - hf_state_dict['k_proj.weight'] = mg_attn_weight[q_dim:-kv_dim, :].clone() - hf_state_dict['v_proj.weight'] = mg_attn_weight[-kv_dim:, :].clone() - if scale_inv is not None: - hf_state_dict['q_proj.weight_scale_inv'] = scale_inv[:q_block, :].clone() - hf_state_dict['k_proj.weight_scale_inv'] = scale_inv[q_block:-kv_block, :].clone() - hf_state_dict['v_proj.weight_scale_inv'] = scale_inv[-kv_block:, :].clone() - del mg_attn_weight - self._set_state_dict(mg_attn, 'linear_proj.weight', hf_state_dict, 'o_proj.weight', to_mcore) - if config.add_qkv_bias and not self._peft_format: - if to_mcore: - linear_qkv_bias = torch.cat([ - hf_state_dict['q_proj.bias'].load(), - hf_state_dict['k_proj.bias'].load(), - hf_state_dict['v_proj.bias'].load(), - ], - dim=0) - self._set_weight(mg_attn.linear_qkv.bias, linear_qkv_bias, 'linear_qkv.bias') - else: - mg_attn_bias, _ = self._get_weight(None if mg_attn is None else mg_attn.linear_qkv.bias.data, - 'linear_qkv.bias') - if mg_attn_bias is not None: - hf_state_dict['q_proj.bias'] = mg_attn_bias[:q_dim].clone() - hf_state_dict['k_proj.bias'] = mg_attn_bias[q_dim:-kv_dim].clone() - hf_state_dict['v_proj.bias'] = mg_attn_bias[-kv_dim:].clone() - self._set_state_dict(mg_attn, 'q_layernorm.weight', hf_state_dict, 'q_norm.weight', to_mcore) - self._set_state_dict(mg_attn, 'k_layernorm.weight', hf_state_dict, 'k_norm.weight', to_mcore) - if to_mcore: - hf_state_dict = {} - else: - hf_state_dict = self._add_prefix(hf_state_dict, hf_prefix) - return hf_state_dict + def _set_qk_layernorm(self, mg_attn, hf_state_dict, to_mcore, **kwargs): + self._set_state_dict(mg_attn, 'q_norm.weight', hf_state_dict, 'q_norm.weight', to_mcore) + self._set_state_dict(mg_attn, 'k_norm.weight', hf_state_dict, 'k_norm.weight', to_mcore) -class OlMoELoader(ModelLoader): +class OLMoELoader(ModelLoader): def get_transformer_layer_spec(self, vp_stage: Optional[int] = None): - return get_olmoe_decoder_block_spec(self.config, vp_stage) + transformer_layer_spec = super().get_transformer_layer_spec(vp_stage) + for layer_spec in transformer_layer_spec.layer_specs: + layer_spec.submodules.self_attention.module = OLMoESelfAttention + return transformer_layer_spec register_model(ModelMeta( ModelType.olmoe, ['olmoe'], bridge_cls=OLMoEBridge, - loader=OlMoELoader, + loader=OLMoELoader, ))