diff --git a/megatron/core/transformer/hyper_connection.py b/megatron/core/transformer/hyper_connection.py index 89d92c88af5..e011f42d929 100644 --- a/megatron/core/transformer/hyper_connection.py +++ b/megatron/core/transformer/hyper_connection.py @@ -195,10 +195,9 @@ def __init__(self, config: TransformerConfig, layer_number: int): ) init_alpha = config.mhc_init_gating_factor - # Learnable scaling factors (Eq. 5 in paper) - self.alpha_pre = nn.Parameter(torch.full((1,), init_alpha)) - self.alpha_post = nn.Parameter(torch.full((1,), init_alpha)) - self.alpha_res = nn.Parameter(torch.full((1,), init_alpha)) + # Learnable scaling factors (Eq. 5 in paper): pre, post, res in one (3,) tensor + # matching HF checkpoint layout (attn_hc.scale / ffn_hc.scale). + self.scale = nn.Parameter(torch.full((3,), init_alpha)) # Static bias terms self.bias = nn.Parameter(torch.zeros(self.n * self.n + 2 * self.n)) @@ -246,9 +245,7 @@ def _init_weights(self) -> None: # (nn.Linear, nn.RMSNorm) whose gradients need to be all-reduced. if self.config.sequence_parallel: setattr(self.mapping_proj.weight, 'sequence_parallel', True) - setattr(self.alpha_pre, 'sequence_parallel', True) - setattr(self.alpha_post, 'sequence_parallel', True) - setattr(self.alpha_res, 'sequence_parallel', True) + setattr(self.scale, 'sequence_parallel', True) setattr(self.bias, 'sequence_parallel', True) def _projection_and_get_norm(self, x: Tensor) -> Tuple[Tensor, Tensor]: @@ -280,9 +277,9 @@ def _compute_h(self, proj: Tensor, r: Tensor) -> Tuple[Tensor, Tensor, Tensor]: """ alpha_ = torch.cat( [ - self.alpha_pre.expand(self.n), - self.alpha_post.expand(self.n), - self.alpha_res.expand(self.n * self.n), + self.scale[0:1].expand(self.n), + self.scale[1:2].expand(self.n), + self.scale[2:3].expand(self.n * self.n), ], dim=-1, ) @@ -319,9 +316,9 @@ def compute_mappings(self, x: Tensor) -> Tuple[Tensor, Tensor, Tensor]: h_pre, h_post, h_res, _ = self._proj_rms_compute_h_op( x_2d, self.mapping_proj.weight, - self.alpha_pre, - self.alpha_post, - self.alpha_res, + self.scale[0:1], + self.scale[1:2], + self.scale[2:3], self.bias, self.n, self.norm_eps, @@ -831,3 +828,37 @@ def compute_optimal_block_size(num_layers: int, num_streams: int) -> int: """ block_size = int(math.sqrt(num_streams * num_layers / (num_streams + 2))) return max(1, block_size) + + +# DSv4 mHC output contraction: n residual streams → single stream. +class HyperConnectionContractModule(MegatronModule): + """Encapsulates the DSv4 mHC output-contraction parameters and forward. + + Parameters are named ``hc_head_fn`` / ``hc_head_base`` / ``hc_head_scale`` + (matching their legacy flat checkpoint keys) so that owners only need to + strip the ``mhc_contract.`` prefix when remapping for backward-compatible + checkpoint loading via ``apply_prefix_mapping``. + """ + + def __init__(self, config: TransformerConfig) -> None: + super().__init__(config) + hc_mult = config.num_residual_streams + hc_dim = config.hidden_size * hc_mult + self.hc_head_fn = nn.Parameter(torch.randn(hc_mult, hc_dim)) + self.hc_head_base = nn.Parameter(torch.zeros(hc_mult)) + self.hc_head_scale = nn.Parameter(torch.ones(1)) + nn.init.xavier_uniform_(self.hc_head_fn) + if config.sequence_parallel: + setattr(self.hc_head_fn, 'sequence_parallel', True) + setattr(self.hc_head_base, 'sequence_parallel', True) + setattr(self.hc_head_scale, 'sequence_parallel', True) + + def forward(self, hidden_states: Tensor) -> Tensor: + return learned_output_contract( + hidden_states, + self.hc_head_fn, + self.hc_head_base, + self.hc_head_scale, + self.config.num_residual_streams, + self.config.layernorm_epsilon, + ) diff --git a/megatron/core/transformer/multi_token_prediction.py b/megatron/core/transformer/multi_token_prediction.py index c0d5ef01e25..33f4753343f 100755 --- a/megatron/core/transformer/multi_token_prediction.py +++ b/megatron/core/transformer/multi_token_prediction.py @@ -28,7 +28,7 @@ inference_all_gather_from_tensor_model_parallel_region, ) from megatron.core.transformer.enums import AttnMaskType, LayerType -from megatron.core.transformer.hyper_connection import learned_output_contract +from megatron.core.transformer.hyper_connection import HyperConnectionContractModule from megatron.core.transformer.module import MegatronModule from megatron.core.transformer.spec_utils import ModuleSpec, build_module from megatron.core.transformer.torch_norm import LayerNormBuilder @@ -1235,16 +1235,7 @@ def __init__( ) if self.mhc_enabled: - hc_mult = self.config.num_residual_streams - hc_dim = self.config.hidden_size * hc_mult - self.hc_head_fn = nn.Parameter(torch.randn(hc_mult, hc_dim)) - self.hc_head_base = nn.Parameter(torch.zeros(hc_mult)) - self.hc_head_scale = nn.Parameter(torch.ones(1)) - nn.init.xavier_uniform_(self.hc_head_fn) - if self.config.sequence_parallel: - setattr(self.hc_head_fn, 'sequence_parallel', True) - setattr(self.hc_head_base, 'sequence_parallel', True) - setattr(self.hc_head_scale, 'sequence_parallel', True) + self.mhc_contract = HyperConnectionContractModule(self.config) self.offload_context = nullcontext() @@ -1453,14 +1444,7 @@ def _postprocess(self, hidden_states: torch.Tensor): """ if self.mhc_enabled: - hidden_states = learned_output_contract( - hidden_states, - self.hc_head_fn, - self.hc_head_base, - self.hc_head_scale, - self.config.num_residual_streams, - self.config.layernorm_epsilon, - ) + hidden_states = self.mhc_contract(hidden_states) # Layer norm before shared head layer. hidden_states = apply_module(self.final_layernorm)(hidden_states) diff --git a/megatron/core/transformer/transformer_block.py b/megatron/core/transformer/transformer_block.py index c3cfa7f6c71..be2e61628ab 100755 --- a/megatron/core/transformer/transformer_block.py +++ b/megatron/core/transformer/transformer_block.py @@ -25,8 +25,8 @@ from megatron.core.tensor_parallel.random import CheckpointManager from megatron.core.transformer.enums import InferenceCudaGraphScope, LayerType from megatron.core.transformer.hyper_connection import ( + HyperConnectionContractModule, HyperConnectionModule, - learned_output_contract, ) from megatron.core.transformer.module import GraphableMegatronModule, MegatronModule from megatron.core.transformer.spec_utils import ModuleSpec, build_module @@ -391,16 +391,7 @@ def build_layer(layer_spec, layer_number): eps=self.config.layernorm_epsilon, ) if self.config.enable_hyper_connections: - hc_mult = self.config.num_residual_streams - hc_dim = self.config.hidden_size * hc_mult - self.hc_head_fn = nn.Parameter(torch.randn(hc_mult, hc_dim)) - self.hc_head_base = nn.Parameter(torch.zeros(hc_mult)) - self.hc_head_scale = nn.Parameter(torch.ones(1)) - nn.init.xavier_uniform_(self.hc_head_fn) - if self.config.sequence_parallel: - setattr(self.hc_head_fn, 'sequence_parallel', True) - setattr(self.hc_head_base, 'sequence_parallel', True) - setattr(self.hc_head_scale, 'sequence_parallel', True) + self.mhc_contract = HyperConnectionContractModule(self.config) else: self.final_layernorm = None # Either this or nn.Identity @@ -956,14 +947,7 @@ def forward( mhc_multistream = hidden_states # DSv4 introduced the new output contraction for mHC. # [s, b, n*C] -> [s, b, C] - hidden_states = learned_output_contract( - hidden_states, - self.hc_head_fn, - self.hc_head_base, - self.hc_head_scale, - self.config.num_residual_streams, - self.config.layernorm_epsilon, - ) + hidden_states = self.mhc_contract(hidden_states) # Final layer norm. if self.final_layernorm is not None: @@ -1080,10 +1064,9 @@ def sharded_state_dict( ) # Save bare parameters/buffers that are direct attributes of this block - # (e.g. hyper-connection learned weights: hc_head_fn, hc_head_base, - # hc_head_scale). The named_children loop above would silently drop - # these since they are not nn.Module children. Mirrors the handling in - # MegatronModule.sharded_state_dict. + # (not nn.Module children — the named_children loop above would silently + # drop them). mhc_contract is now a proper child module and is handled + # above; this catches any other bare params on the block itself. local_state_dict: dict = {} self._save_to_state_dict(local_state_dict, '', keep_vars=True) if local_state_dict: diff --git a/tests/unit_tests/transformer/test_multi_token_prediction.py b/tests/unit_tests/transformer/test_multi_token_prediction.py index 849d8a4d42c..c1e01e9832b 100644 --- a/tests/unit_tests/transformer/test_multi_token_prediction.py +++ b/tests/unit_tests/transformer/test_multi_token_prediction.py @@ -1765,11 +1765,11 @@ def test_mtp_constructor_with_mhc(self, tp): assert layer.eh_proj is None assert layer.e_proj.weight.shape == (h // tp, h) assert layer.h_proj.weight.shape == (h // tp, h) - assert layer.hc_head_fn.shape == (n, n * h) - assert layer.hc_head_base.shape == (n,) - assert layer.hc_head_scale.shape == (1,) + assert layer.mhc_contract.hc_head_fn.shape == (n, n * h) + assert layer.mhc_contract.hc_head_base.shape == (n,) + assert layer.mhc_contract.hc_head_scale.shape == (1,) if tp > 1: - assert getattr(layer.hc_head_fn, 'sequence_parallel', False) + assert getattr(layer.mhc_contract.hc_head_fn, 'sequence_parallel', False) def test_transformer_block_returns_tuple(self): """With mHC+MTP the block returns (contracted, multistream); without MTP just a tensor.""" diff --git a/tests/unit_tests/transformer/test_transformer_layer.py b/tests/unit_tests/transformer/test_transformer_layer.py index b7dbe27e71a..32490d11d8f 100644 --- a/tests/unit_tests/transformer/test_transformer_layer.py +++ b/tests/unit_tests/transformer/test_transformer_layer.py @@ -828,7 +828,7 @@ def test_submodules_under_cudagraphs_includes_hyper_connection(self): ) assert hc_modules_found, ( "_get_submodules_under_cudagraphs does not include HyperConnectionModule. " - "Parameters like mapping_proj, alpha_pre/post/res will not be updated " + "Parameters like mapping_proj, scale, bias will not be updated " "during CUDA graph replay." )