Skip to content

standalone hypercontraction module #5669

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Open
liuzhenhai93 wants to merge 1 commit into
base: dev
Choose a base branch
from
Open

standalone hypercontraction module #5669

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
1 commit
Select commit
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
57 changes: 44 additions & 13 deletions megatron/core/transformer/hyper_connection.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -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))
Comment on lines -198 to +200

@cspades cspades Jul 8, 2026

Copy link
Copy Markdown
Member

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

This will affect backwards compatibility with earlier checkpoints right?

Can you handle the HF parameter format online during checkpoint loading?


# Static bias terms
self.bias = nn.Parameter(torch.zeros(self.n * self.n + 2 * self.n))
Expand Down Expand Up @@ -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]:
Expand Down Expand Up @@ -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,
)
Expand Down Expand Up @@ -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,
Expand Down Expand Up @@ -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``.
Comment on lines +837 to +840

@claude claude Bot Jul 6, 2026

Copy link
Copy Markdown
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

[SUGGESTION Other] This docstring describes backward-compatible checkpoint loading "via apply_prefix_mapping" as if it were an available/used mechanism, but no callsite in this PR (neither transformer_block.py nor multi_token_prediction.py) actually performs that remap. A docstring that describes behavior the code doesn't implement is worse than none — either wire up the apply_prefix_mapping call at the owners (see the block/MTP comments) or reword this to state plainly that owners must add the remap themselves for BC, and that this module alone does not preserve the legacy flat keys.

"""

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,
)
22 changes: 3 additions & 19 deletions megatron/core/transformer/multi_token_prediction.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -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
Expand Down Expand Up @@ -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)

@claude claude Bot Jul 6, 2026

Copy link
Copy Markdown
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

[IMPORTANT Compatibility] Same checkpoint-key rename issue as in transformer_block.py. Previously the MTP layer saved hc_head_fn/hc_head_base/hc_head_scale as bare params (keys ...mtp.layers.{i}.hc_head_fn); now they are under self.mhc_contract, producing ...mtp.layers.{i}.mhc_contract.hc_head_fn.

This layer already has a sharded_state_dict override (line 1763) that calls apply_prefix_mapping for the mtp_model_layer./transformer_layer. remap — extend it to also strip the mhc_contract. prefix when self.mhc_enabled, so existing mHC+MTP checkpoints keep loading. Otherwise resume from an old checkpoint silently loses these learned params.

@cspades cspades Jul 8, 2026

Copy link
Copy Markdown
Member

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Yup, pretty much my concern as well, we need to make sure that checkpointing works as a requirement to performing this refactor.


self.offload_context = nullcontext()

Expand Down Expand Up @@ -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)
Expand Down
29 changes: 6 additions & 23 deletions megatron/core/transformer/transformer_block.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -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
Expand Down Expand Up @@ -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

Expand Down Expand Up @@ -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:
Expand Down Expand Up @@ -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.
Comment on lines 1066 to +1069

@claude claude Bot Jul 6, 2026

Copy link
Copy Markdown
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

[IMPORTANT Compatibility] This refactor silently renames the checkpoint keys for the mHC contraction params without any migration path.

Before this PR, hc_head_fn/hc_head_base/hc_head_scale were bare nn.Parameter attributes of the block, saved by the _save_to_state_dict(...) + make_sharded_tensors_for_checkpoint(local_state_dict, prefix, ...) path below — producing keys like decoder.hc_head_fn. Now they live under self.mhc_contract, so the named_children() loop (lines 1053-1064) saves them as decoder.mhc_contract.hc_head_fn, etc.

Impact: existing checkpoints from PR #4518 (already on dev) that contain ...hc_head_fn will no longer load into the new code — the keys won't match, causing a silent divergence or a load error on resume.

Fix: wire the remap the new module's own docstring promises. In this block's sharded_state_dict, after building sharded_state_dict, strip the new prefix for backward-compat, e.g.:

from megatron.core.dist_checkpointing.utils import apply_prefix_mapping
if self.config.enable_hyper_connections:
    apply_prefix_mapping(
        sharded_state_dict,
        {f'{prefix}mhc_contract.hc_head_fn': f'{prefix}hc_head_fn',
         f'{prefix}mhc_contract.hc_head_base': f'{prefix}hc_head_base',
         f'{prefix}mhc_contract.hc_head_scale': f'{prefix}hc_head_scale'},
    )

(or otherwise document that this breaks existing mHC checkpoints). As written, the docstring claims a migration that no callsite actually performs.

local_state_dict: dict = {}
self._save_to_state_dict(local_state_dict, '', keep_vars=True)
if local_state_dict:
Expand Down
8 changes: 4 additions & 4 deletions tests/unit_tests/transformer/test_multi_token_prediction.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -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."""
Expand Down
2 changes: 1 addition & 1 deletion tests/unit_tests/transformer/test_transformer_layer.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -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."
)

Expand Down