NVIDIA / NVIDIA/Megatron-LM

Allow disabling grouped GEMM for Mamba/hybrid

Open
#4,106 0 comments 0 reactions 0 assignees View on GitHub
enhancement
Dominant language
Python
Stars
17.9k
Forks
4.5k
Avg merge
4d 6h
Merged PRs (30d)
271

Description

Originally posted here:
https://github.com/NVIDIA/Megatron-LM/pull/4099/changes#r3026195753
cc @kevalmorabia97

**Is your feature request related to a problem? Please describe.**
The Mamba/hybrid model currently does not support setting the `moe_grouped_gemm` via `config` in `megatron/core/models/mamba/mamba_layer_specs.py`/`megatron/core/models/hybrid/hybrid_layer_specs.py`, unlike the GPT model. ModelOpt pruning does not support grouped GEMM, so this would enable easier support without having to modify the code itself.

**Describe the solution you'd like**
(Refers to `megatron/core/models/mamba/mamba_layer_specs.py`, which may have been renamed to `megatron/core/models/hybrid/hybrid_layer_specs.py` in the meantime.)

> Can we add a convenience function `get_te_mamba_stack_spec(moe_grouped_gemm: bool = False)` so we can configure whether to use moe grouped gemm or not? This is similar to gpt model spec function - this is useful for ModelOpt pruning where we dont support grouped gemm

Contributor guide

Open the contributing guide

Assessment

This issue has not been assessed yet.

Get new issues in your inbox

A short digest of beginner-friendly GitHub issues.