seperate_hypercontraction_module#5669
Conversation
|
/claude strict-review |
Strict Review SummaryFindings: CRITICAL: 0, IMPORTANT: 2, SUGGESTION: 1 This PR is a clean, behavior-preserving refactor of the DSv4 mHC output-contraction params into a dedicated Most impactful findings
Risk assessmentMedium. The compute path is safe, but the checkpoint-key change is a real backward-compatibility regression for anyone with existing mHC checkpoints. Should not merge until either the prefix mapping is added or the break is explicitly acknowledged/documented. A |
![claude[bot]](https://avatars.githubusercontent.com/in/1236702?s=60&v=4){kind=small}
| # 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. |
There was a problem hiding this comment.
[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.
| 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) |
There was a problem hiding this comment.
[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.
| 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``. |
There was a problem hiding this comment.
[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.
3 participants
What does this PR do?
refactor hyper connection contraction to a module
Issue tracking
For PRs from open-source community contributors:
Linked issue:
#5670
Contribution process
Pre-checks
Code review
Feel free to message or comment @NVIDIA/mcore-oncall to help accelerate your merge into main. The less complex your PR is, the faster it will be approved and merged!
All PRs start as draft. If you open a non-draft PR, it will be automatically converted to draft.
Step 1: Mark PR as "Ready for Review"
.github/CODEOWNERS.Final Review might get declined if these requirements are not fulfilled.
Step 2: Final Review
For PRs that change
megatron/core, once all expert reviewers have approved, theFinal Reviewlabel is applied automatically and final reviewers are assigned.For PRs outside
megatron/core, this step is skipped.Step 3: Approved
Once all required reviewers have approved, the
Approvedlabel is applied automatically.Merge
Any member of mcore-engineers will be able to merge your PR.