diff --git a/src/mcore_bridge/bridge/gpt_bridge.py b/src/mcore_bridge/bridge/gpt_bridge.py index 41caf7c..4894559 100644 --- a/src/mcore_bridge/bridge/gpt_bridge.py +++ b/src/mcore_bridge/bridge/gpt_bridge.py @@ -51,6 +51,10 @@ class GPTBridge: additional_dim1_keys = set() _support_hf_grouped_lora = True + @property + def use_transformer_engine(self): + return self.config.transformer_impl == 'transformer_engine' + def __init__(self, config: ModelConfig): self.config = config self._disable_tqdm = False @@ -1638,8 +1642,9 @@ def _set_layer_attn(self, mg_layer, hf_state_dict, layer_idx: int, to_mcore: boo else: hf_state_dict.update( self._set_attn_state(mg_attn, hf_state_dict, f'{self.hf_attn_prefix}.', layer_idx, to_mcore)) - self._set_state_dict(mg_layer, 'self_attention.linear_qkv.layer_norm_weight', hf_state_dict, - self.hf_input_layernorm_key, to_mcore) + mg_key = ('self_attention.linear_qkv.layer_norm_weight' + if self.use_transformer_engine else 'input_layernorm.weight') + self._set_state_dict(mg_layer, mg_key, hf_state_dict, self.hf_input_layernorm_key, to_mcore) return hf_state_dict def _set_layer_mlp(self, mg_layer, hf_state_dict, layer_idx: int, to_mcore: bool, is_mtp: bool = False): @@ -1658,8 +1663,8 @@ def _set_layer_mlp(self, mg_layer, hf_state_dict, layer_idx: int, to_mcore: bool else: hf_state_dict.update( self._set_mlp_state(mg_mlp, hf_state_dict, f'{self.hf_mlp_prefix}.', layer_idx, to_mcore)) - self._set_state_dict(mg_layer, 'mlp.linear_fc1.layer_norm_weight', hf_state_dict, - self.hf_post_attention_layernorm_key, to_mcore) + mg_key = 'mlp.linear_fc1.layer_norm_weight' if self.use_transformer_engine else 'pre_mlp_layernorm.weight' + self._set_state_dict(mg_layer, mg_key, hf_state_dict, self.hf_post_attention_layernorm_key, to_mcore) return hf_state_dict def _set_hyper_connection(self, mg_layer, hf_state_dict, layer_idx, to_mcore): diff --git a/src/mcore_bridge/model/mm_gpts/gemma4.py b/src/mcore_bridge/model/mm_gpts/gemma4.py index 21624ee..2288794 100644 --- a/src/mcore_bridge/model/mm_gpts/gemma4.py +++ b/src/mcore_bridge/model/mm_gpts/gemma4.py @@ -468,8 +468,8 @@ def _set_router(self, mg_mlp, hf_state_dict, to_mcore, **kwargs): def _set_layer_mlp(self, mg_layer, hf_state_dict, layer_idx: int, to_mcore: bool, is_mtp: bool = False): mg_mlp = None if mg_layer is None else mg_layer.mlp hf_state_dict.update(self._set_mlp_state(mg_mlp, hf_state_dict, f'{self.hf_mlp_prefix}.', layer_idx, to_mcore)) - self._set_state_dict(mg_layer, 'mlp.linear_fc1.layer_norm_weight', hf_state_dict, - 'pre_feedforward_layernorm.weight', to_mcore) + mg_key = 'mlp.linear_fc1.layer_norm_weight' if self.use_transformer_engine else 'pre_mlp_layernorm.weight' + self._set_state_dict(mg_layer, mg_key, hf_state_dict, 'pre_feedforward_layernorm.weight', to_mcore) if self.text_config.enable_moe_block: mg_experts = None if mg_layer is None else mg_layer.experts_mlp hf_state_dict.update(self._set_moe_state(mg_experts, hf_state_dict, '', layer_idx, to_mcore, is_mtp=is_mtp)) @@ -840,14 +840,20 @@ def get_transformer_layer_spec(self, vp_stage: Optional[int] = None): num_moe_experts = self.config.num_moe_experts self.config.num_moe_experts = None layer_specs = get_gpt_decoder_block_spec( - self.config, use_transformer_engine=True, normalization=self.config.normalization, vp_stage=vp_stage) + self.config, + use_transformer_engine=self.use_transformer_engine, + normalization=self.config.normalization, + vp_stage=vp_stage) for layer_spec in layer_specs.layer_specs: layer_spec.submodules.self_attention.module = Gemma4SelfAttention self._set_mlp_spec(layer_spec.submodules, Gemma4MLP) if num_moe_experts is not None: self.config.num_moe_experts = num_moe_experts moe_layer_specs = get_gpt_decoder_block_spec( - self.config, use_transformer_engine=True, normalization=self.config.normalization, vp_stage=vp_stage) + self.config, + use_transformer_engine=self.use_transformer_engine, + normalization=self.config.normalization, + vp_stage=vp_stage) for layer_spec, moe_layer_spec in zip(layer_specs.layer_specs, moe_layer_specs.layer_specs): layer_spec.submodules.experts_mlp = moe_layer_spec.submodules.mlp self._set_mlp_spec(layer_spec.submodules, Gemma4MoELayer, mlp_key='experts_mlp') diff --git a/src/mcore_bridge/model/register.py b/src/mcore_bridge/model/register.py index e0d6541..414a2de 100644 --- a/src/mcore_bridge/model/register.py +++ b/src/mcore_bridge/model/register.py @@ -9,11 +9,13 @@ from megatron.core.extensions.transformer_engine import TEGroupedLinear, TELayerNormColumnParallelLinear, TELinear from megatron.core.models.gpt import gpt_model from megatron.core.models.gpt.gpt_layer_specs import get_gpt_decoder_block_spec, get_gpt_mtp_block_spec +from megatron.core.transformer.dot_product_attention import DotProductAttention from megatron.core.transformer.moe.router import TopKRouter as McoreTopKRouter from megatron.core.transformer.multi_latent_attention import MLASelfAttention as McoreMLASelfAttention from megatron.core.transformer.transformer_layer import TransformerLayer as McoreTransformerLayer from packaging import version from torch import nn +from transformers.utils import is_torch_npu_available from typing import TYPE_CHECKING, List, Optional, Type, Union from mcore_bridge.bridge import GPTBridge @@ -82,6 +84,10 @@ def __init__(self, config: ModelConfig): if self.model_cls is None: self.model_cls = MultimodalGPTModel if config.is_multimodal else GPTModel + @property + def use_transformer_engine(self): + return self.config.transformer_impl == 'transformer_engine' + def _set_mlp_spec(self, layer_submodules, mlp_module, mlp_key='mlp'): mlp_spec = getattr(layer_submodules, mlp_key) if isinstance(mlp_spec, partial): @@ -121,15 +127,22 @@ def _deepcopy_layer_spec(self, transformer_layer_spec): for i, layer_spec in enumerate(transformer_layer_spec.layer_specs): transformer_layer_spec.layer_specs[i] = deepcopy(layer_spec) + def _replace_unfused_attention(self, transformer_layer_spec): + if not is_torch_npu_available() or self.config.attention_backend.name != 'unfused': + return + for layer_spec in transformer_layer_spec.layer_specs: + layer_spec.submodules.self_attention.submodules.core_attention = DotProductAttention + def get_transformer_layer_spec(self, vp_stage: Optional[int] = None): with self._patch_experimental_attention_variant(): transformer_layer_spec = get_gpt_decoder_block_spec( self.config, - use_transformer_engine=True, + use_transformer_engine=self.use_transformer_engine, normalization=self.config.normalization, qk_l2_norm=self.config.qk_l2_norm, vp_stage=vp_stage) self._deepcopy_layer_spec(transformer_layer_spec) + self._replace_unfused_attention(transformer_layer_spec) if self.config.experimental_attention_variant == 'dsa': for layer_spec in transformer_layer_spec.layer_specs: self._replace_spec_dsa(layer_spec) @@ -137,7 +150,7 @@ def get_transformer_layer_spec(self, vp_stage: Optional[int] = None): def get_mtp_block_spec(self, transformer_layer_spec, vp_stage: Optional[int] = None): mtp_block_spec = get_gpt_mtp_block_spec( - self.config, transformer_layer_spec, use_transformer_engine=True, vp_stage=vp_stage) + self.config, transformer_layer_spec, use_transformer_engine=self.use_transformer_engine, vp_stage=vp_stage) if mtp_block_spec is not None: for layer_spec in mtp_block_spec.layer_specs: layer_spec.module = MultiTokenPredictionLayer diff --git a/src/mcore_bridge/tuners/lora.py b/src/mcore_bridge/tuners/lora.py index 8b15093..890d808 100644 --- a/src/mcore_bridge/tuners/lora.py +++ b/src/mcore_bridge/tuners/lora.py @@ -15,6 +15,7 @@ TERowParallelGroupedLinear, TERowParallelLinear) from megatron.core.models.common.embeddings.language_model_embedding import LanguageModelEmbedding from megatron.core.parallel_state import get_expert_tensor_parallel_world_size, get_tensor_model_parallel_world_size +from megatron.core.tensor_parallel.layers import ColumnParallelLinear, RowParallelLinear from megatron.core.tensor_parallel.random import get_cuda_rng_tracker, get_expert_parallel_rng_tracker_name from megatron.core.transformer.mlp import apply_swiglu_sharded_factory from megatron.core.transformer.module import MegatronModule @@ -117,6 +118,8 @@ def __init__( raise ValueError(f'{self.__class__.__name__} does not support DoRA yet, please set it to False') self.is_parallel_a = isinstance(base_layer, (TERowParallelLinear, TERowParallelGroupedLinear)) + self.is_local_linear = isinstance(base_layer, (ColumnParallelLinear, RowParallelLinear)) + self.is_parallel_a = self.is_parallel_a or isinstance(base_layer, RowParallelLinear) self.is_grouped = isinstance(base_layer, TEGroupedLinear) self.fan_in_fan_out = fan_in_fan_out self._active_adapter = adapter_name @@ -180,7 +183,7 @@ def update_layer(self, adapter_name, r, *, lora_alpha, **kwargs): lora_a = _build_local_te_linear(router_shape[1], r, lora_bias, **kwargs) lora_b = _build_local_te_linear(r, router_shape[0], lora_bias, **kwargs) elif self.is_parallel_a: - in_features = self.in_features * self.tp_size + in_features = self.in_features if self.is_local_linear else self.in_features * self.tp_size if self.is_grouped: if is_torch_npu_available(): lora_a = NpuGroupedLoraLinear( @@ -214,15 +217,25 @@ def update_layer(self, adapter_name, r, *, lora_alpha, **kwargs): **kwargs, ) else: - lora_a = TERowParallelLinear( - input_size=in_features, - output_size=r, - bias=False, - input_is_parallel=True, - **kwargs, - ) - lora_b = _build_local_te_linear(r, self.out_features, lora_bias, **kwargs) - lora_a.parallel_mode = self.base_layer.parallel_mode # fix moe_shared_expert_overlap + if self.is_local_linear: + lora_a = RowParallelLinear( + input_size=in_features, + output_size=r, + bias=False, + input_is_parallel=True, + **kwargs, + ) + lora_b = nn.Linear(r, self.out_features, bias=lora_bias) + else: + lora_a = TERowParallelLinear( + input_size=in_features, + output_size=r, + bias=False, + input_is_parallel=True, + **kwargs, + ) + lora_b = _build_local_te_linear(r, self.out_features, lora_bias, **kwargs) + lora_a.parallel_mode = self.base_layer.parallel_mode # fix moe_shared_expert_overlap else: if is_torch_npu_available(): out_features = self.out_features @@ -260,15 +273,25 @@ def update_layer(self, adapter_name, r, *, lora_alpha, **kwargs): **kwargs, ) else: - lora_a = _build_local_te_linear(self.in_features, r, lora_bias, **kwargs) - lora_b = TEColumnParallelLinear( - input_size=r, - output_size=out_features, - bias=lora_bias, - gather_output=False, - **kwargs, - ) - lora_b.parallel_mode = self.base_layer.parallel_mode # fix moe_shared_expert_overlap + if self.is_local_linear: + lora_a = nn.Linear(self.in_features, r, bias=lora_bias) + lora_b = ColumnParallelLinear( + input_size=r, + output_size=out_features, + bias=lora_bias, + gather_output=False, + **kwargs, + ) + else: + lora_a = _build_local_te_linear(self.in_features, r, lora_bias, **kwargs) + lora_b = TEColumnParallelLinear( + input_size=r, + output_size=out_features, + bias=lora_bias, + gather_output=False, + **kwargs, + ) + lora_b.parallel_mode = self.base_layer.parallel_mode # fix moe_shared_expert_overlap for lora in [lora_a, lora_b]: # When parallel_mode is set to None by moe_shared_expert_overlap, # disable UB comm overlap; the corresponding collectives are driven @@ -412,7 +435,7 @@ def forward(self, x: torch.Tensor, *args: Any, **kwargs: Any): f'Got base_layer type: {type(self.base_layer)}. ') else: (result, x), bias = self.base_layer(x, *args, **kwargs) - elif isinstance(self.base_layer, (TELinear, TEGroupedLinear)): + elif isinstance(self.base_layer, (TELinear, TEGroupedLinear, ColumnParallelLinear, RowParallelLinear)): result, bias = self.base_layer(x, *args, **kwargs) elif isinstance(self.base_layer, TopKRouter): with self._patch_router_gating(): diff --git a/src/mcore_bridge/tuners/patcher.py b/src/mcore_bridge/tuners/patcher.py index f9cae8c..4eda69a 100644 --- a/src/mcore_bridge/tuners/patcher.py +++ b/src/mcore_bridge/tuners/patcher.py @@ -1,5 +1,6 @@ # Copyright (c) ModelScope Contributors. All rights reserved. from megatron.core.extensions.transformer_engine import TEGroupedLinear, TELayerNormColumnParallelLinear, TELinear +from megatron.core.tensor_parallel.layers import ColumnParallelLinear, RowParallelLinear from megatron.core.transformer.module import MegatronModule from megatron.core.transformer.moe.router import TopKRouter from peft import LoraModel @@ -27,7 +28,8 @@ def dispatch_megatron( else: target_base_layer = target - linear_cls = (TELayerNormColumnParallelLinear, TELinear, TEGroupedLinear, TopKRouter) + linear_cls = (TELayerNormColumnParallelLinear, TELinear, TEGroupedLinear, ColumnParallelLinear, RowParallelLinear, + TopKRouter) if isinstance(target_base_layer, linear_cls): new_module = LoraParallelLinear(base_layer=target, adapter_name=adapter_name, **kwargs) diff --git a/tests/test_model_register.py b/tests/test_model_register.py new file mode 100644 index 0000000..1e4aa4f --- /dev/null +++ b/tests/test_model_register.py @@ -0,0 +1,96 @@ +import unittest +from types import SimpleNamespace +from unittest.mock import patch + +from mcore_bridge.bridge.gpt_bridge import GPTBridge +from mcore_bridge.model.register import DotProductAttention, ModelLoader + + +class TestModelLoader(unittest.TestCase): + + @staticmethod + def _get_loader(transformer_impl): + loader = ModelLoader.__new__(ModelLoader) + loader.config = SimpleNamespace( + transformer_impl=transformer_impl, + experimental_attention_variant=None, + normalization='RMSNorm', + qk_l2_norm=False, + attention_backend=SimpleNamespace(name='flash'), + ) + return loader + + def test_transformer_impl_controls_decoder_layer_spec(self): + for transformer_impl, expected in [('local', False), ('transformer_engine', True)]: + with self.subTest(transformer_impl=transformer_impl): + loader = self._get_loader(transformer_impl) + with patch('mcore_bridge.model.register.get_gpt_decoder_block_spec') as get_spec: + get_spec.return_value = SimpleNamespace(layer_specs=[]) + loader.get_transformer_layer_spec() + self.assertEqual(get_spec.call_args.kwargs['use_transformer_engine'], expected) + + @patch('mcore_bridge.model.register.is_torch_npu_available', return_value=True) + def test_unfused_attention_uses_mcore_dot_product_on_npu(self, _): + loader = self._get_loader('transformer_engine') + loader.config.attention_backend.name = 'unfused' + core_attention = object() + transformer_layer_spec = SimpleNamespace(layer_specs=[ + SimpleNamespace( + submodules=SimpleNamespace( + self_attention=SimpleNamespace(submodules=SimpleNamespace(core_attention=core_attention)))) + ]) + + loader._replace_unfused_attention(transformer_layer_spec) + + self.assertIs(transformer_layer_spec.layer_specs[0].submodules.self_attention.submodules.core_attention, + DotProductAttention) + + def test_transformer_impl_controls_mtp_layer_spec(self): + for transformer_impl, expected in [('local', False), ('transformer_engine', True)]: + with self.subTest(transformer_impl=transformer_impl): + loader = self._get_loader(transformer_impl) + with patch('mcore_bridge.model.register.get_gpt_mtp_block_spec') as get_spec: + get_spec.return_value = None + loader.get_mtp_block_spec(SimpleNamespace()) + self.assertEqual(get_spec.call_args.kwargs['use_transformer_engine'], expected) + + +class TestGPTBridge(unittest.TestCase): + + @staticmethod + def _get_bridge(transformer_impl): + bridge = GPTBridge.__new__(GPTBridge) + bridge.config = SimpleNamespace(transformer_impl=transformer_impl, multi_latent_attention=False) + return bridge + + def test_transformer_impl_controls_attention_layernorm_mapping(self): + expected_keys = { + 'local': 'input_layernorm.weight', + 'transformer_engine': 'self_attention.linear_qkv.layer_norm_weight', + } + layer = SimpleNamespace(self_attention=SimpleNamespace()) + for transformer_impl, expected_key in expected_keys.items(): + with self.subTest(transformer_impl=transformer_impl): + bridge = self._get_bridge(transformer_impl) + with patch.object(bridge, '_set_attn_state', return_value={}), \ + patch.object(bridge, '_set_state_dict') as set_state_dict: + bridge._set_layer_attn(layer, {}, 0, True) + self.assertEqual(set_state_dict.call_args.args[1], expected_key) + + def test_transformer_impl_controls_mlp_layernorm_mapping(self): + expected_keys = { + 'local': 'pre_mlp_layernorm.weight', + 'transformer_engine': 'mlp.linear_fc1.layer_norm_weight', + } + layer = SimpleNamespace(mlp=SimpleNamespace()) + for transformer_impl, expected_key in expected_keys.items(): + with self.subTest(transformer_impl=transformer_impl): + bridge = self._get_bridge(transformer_impl) + with patch.object(bridge, '_set_mlp_state', return_value={}), \ + patch.object(bridge, '_set_state_dict') as set_state_dict: + bridge._set_layer_mlp(layer, {}, 0, True) + self.assertEqual(set_state_dict.call_args.args[1], expected_key) + + +if __name__ == '__main__': + unittest.main() diff --git a/tests/test_tuners.py b/tests/test_tuners.py new file mode 100644 index 0000000..7ac9327 --- /dev/null +++ b/tests/test_tuners.py @@ -0,0 +1,22 @@ +import unittest +from unittest.mock import patch + +from megatron.core.tensor_parallel.layers import ColumnParallelLinear, RowParallelLinear + +from mcore_bridge.tuners.patcher import dispatch_megatron + + +class TestMegatronLoraDispatch(unittest.TestCase): + + def test_dispatches_local_parallel_linears(self): + for linear_cls in (ColumnParallelLinear, RowParallelLinear): + with self.subTest(linear_cls=linear_cls.__name__): + target = linear_cls.__new__(linear_cls) + with patch('mcore_bridge.tuners.patcher.LoraParallelLinear', return_value='adapter') as lora_cls: + result = dispatch_megatron(target, 'default') + self.assertEqual(result, 'adapter') + self.assertIs(lora_cls.call_args.kwargs['base_layer'], target) + + +if __name__ == '__main__': + unittest.main()