NVIDIA / NVIDIA/Megatron-LM

Is there a hard blocker preventing Muon to be used with gated deltanet and gated attention?

Open
#2,885 2 comments 0 reactions 1 assignee Claimed by @janEbert View on GitHub
community-request question waiting-on-maintainers
Dominant language
Python
Stars
17.9k
Forks
4.5k
Avg merge
4d 3h
Merged PRs (30d)
272

Description

Hi @FDecaYed I would like to consult you on using Muon with linear attention and gated attention.

In https://github.com/NVIDIA/Megatron-LM/commit/5880674a396ed3489d185afaefb1773a6481fe18#diff-42fa19ec8893eabf951a5bb21edfc0dfe7c9c8949d5087ab133f5897ea0e3213 it is mentioned in the arguments line 1203 that
```
assert args.linear_attention_type is None, "Muon optimizer does not support linear attention type for now."
assert not args.attention_output_gate, "Muon optimizer does not support attention output gate for now."
```

However, is there a hard incompatibility, or is it just not supported yet?

Could they simply be routed to Adam while the rest of the layers continue to use Muon?

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.