Bridges for Mamba-based models #554
Open
Conversation
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Signed-off-by: Maanu Grover <[email protected]>
Signed-off-by: Maanu Grover <[email protected]>
Signed-off-by: Maanu Grover <[email protected]>
ananthsub
reviewed
Sep 4, 2025
Comment on lines
+30
to
+49
| class PrunedVocabMapping(AutoMapping): | ||
| """ | ||
| Smart mapping like AutoMapping that additionally prunes vocab padding. | ||
|
|
||
| Intended for embedding and output layers. | ||
| """ | ||
|
|
||
| def megatron_to_hf( | ||
| self, | ||
| megatron_weights: Optional[torch.Tensor], | ||
| megatron_module: Optional[nn.Module], | ||
| ) -> dict[str, torch.Tensor]: | ||
| """Prune padding from weight in vocab size dimension, if vocab size is accessible.""" | ||
| mapping = super().megatron_to_hf(megatron_weights, megatron_module) | ||
|
|
||
| if megatron_module is not None: | ||
| weight = mapping[str(self.hf_param)] | ||
| mapping[str(self.hf_param)] = weight[: megatron_module.vocab_size, :] | ||
|
|
||
| return mapping |
There was a problem hiding this comment.
should this be a generic transform made available here? https://github.com/NVIDIA-NeMo/Megatron-Bridge/blob/main/src/megatron/bridge/models/conversion/param_mapping.py
There was a problem hiding this comment.
Was discussing this with @yaoyu-33 . In Nemo 2.0, this pruning is only done for the NemotronH exporter: https://github.com/NVIDIA-NeMo/NeMo/blob/main/nemo/collections/llm/gpt/model/ssm.py#L817-L850
I'm not sure why this wasn't needed for other models.
I think @JRD971000 mentioned that the NemotronH checkpoints were saved with vocab padding so this was a choice made when writing exporters specifically for NemotronH.
Signed-off-by: Maanu Grover <[email protected]>
Signed-off-by: Maanu Grover <[email protected]>
|
Needs #440 to be merged first, for HF remote model support. |
Signed-off-by: Maanu Grover <[email protected]>
Signed-off-by: Maanu Grover <[email protected]>
Signed-off-by: Maanu Grover <[email protected]>
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
Add this suggestion to a batch that can be applied as a single commit.
This suggestion is invalid because no changes were made to the code.
Suggestions cannot be applied while the pull request is closed.
Suggestions cannot be applied while viewing a subset of changes.
Only one suggestion per line can be applied in a batch.
Add this suggestion to a batch that can be applied as a single commit.
Applying suggestions on deleted lines is not supported.
You must change the existing code in this line in order to create a valid suggestion.
Outdated suggestions cannot be applied.
This suggestion has been applied or marked resolved.
Suggestions cannot be applied from pending reviews.
Suggestions cannot be applied on multi-line comments.
Suggestions cannot be applied while the pull request is queued to merge.
Suggestion cannot be applied right now. Please check back later.
Add two bridges for Mamba-based models:
MambaBridge which supports any
MambaForCausalLMincluding:NemotronHBridge which supports any
NemotronHForCausalLMincluding: