Skip to content

fix(optim): Muon routing for depthwise conv1d and grouped linear weights #4120

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

Closed
faresobeid wants to merge 1 commit into from
Closed

fix(optim): Muon routing for depthwise conv1d and grouped linear weights #4120

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
13 changes: 13 additions & 0 deletions src/prime_rl/trainer/optim/__init__.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -7,6 +7,7 @@

from prime_rl.configs.trainer import OffloadMode, OptimizerConfig
from prime_rl.trainer.models.fusions import get_model_packed_parameters
from prime_rl.trainer.models.layers.grouped_linear import GroupedLinear
from prime_rl.trainer.optim.base import OffloadOptimizer as OffloadOptimizer
from prime_rl.trainer.optim.base import OptimizerLike
from prime_rl.trainer.optim.offload import (
Expand Down Expand Up @@ -140,9 +141,17 @@ def _create_muon_optimizer(
model: nn.Module | None,
lr: float | None = None,
) -> Optimizer:
# Depthwise convolutions (gated delta net, Mamba, PLE) hold per-channel filters, not matrices: Muon would
# orthogonalize each channel's taps on their own, so they go to AdamW.
conv_weights = set()
if model is not None:
conv_weights = {f"{name}.weight" for name, module in model.named_modules() if isinstance(module, nn.Conv1d)}

def muon_enabled(n, p):
if p.ndim < 2:
return False
if n in conv_weights:
return False
if "lm_head" in n:
return False
if "embed_tokens" in n:
Expand Down Expand Up @@ -218,6 +227,10 @@ def muon_enabled(n, p):
partitions = packed_info.spec.muon_matrix_partitions(packed_info.parameter)
if partitions is not None and packed_info.parameter in muon_params:
matrix_partitions[packed_info.parameter] = partitions
# A grouped linear is n_groups independent matrices stacked along the output rows.
for module in model.modules():
if isinstance(module, GroupedLinear) and module.weight in muon_params:
matrix_partitions[module.weight] = (module.out_features // module.n_groups,) * module.n_groups

optimizer = Muon(
params=param_groups,
Expand Down
Loading