Hacktoberfest 2026: los issues que los mantenedores marcaron para octubre, abiertos y aptos para principiantes. Explorar issues de Hacktoberfest

[Feature] Per-Head Muon: per-head Newton-Schulz orthogonalization for attention projections

Abierto
  • ¿Algo no está bien?
#8,367 5 comentarios 4 reacciones 0 asignados Ver en GitHub

Los mantenedores suelen responder en 3 días

@0z5a ya está trabajando en esto.

Desde el 18/9/2026.

Evaluación

Dificultad
5/5
Tiempo estimado
Más de una semana
Aptitud para principiantes
45/100
Tipo de issue
Nueva funcionalidad
Claridad
Bastante claro
Estado de actividad
Activo
Stack tecnológico
python
Área
distributed-systems, machine-learning, performance

Línea de trabajo

Empieza leyendo deepspeed/runtime/zero/muon/original_muon.py y muon_optimizer.py para entender muon_update y sus rutas de Newton–Schulz; después, inspecciona el punto de llamada de ZeRO en deepspeed/runtime/zero/stage_1_and_2.py. Traza cómo se propagan los metadatos existentes de use_muon y la estructura de atención de AutoTP. Se considera terminado cuando un modo opt-in por cabeza funciona a través de ZeRO y está cubierto por pruebas de ortogonalización sintética y de convergencia con modelos pequeños.

Escrito por el modelo de indexación a partir del texto del issue.

Descripción

enhancement training

Motivation

DeepSpeed now supports the Muon optimizer across ZeRO stages 1/2/3. In 2026, two frontier labs independently converged on the same refinement for attention projections: orthogonalizing per attention head rather than over the full projection matrix.

  • Kimi K3 (arXiv:2607.24653, §2.5, "Per-Head Muon"): partitions the momentum of Q/K/V projections along the head dimension and orthogonalizes each head's block separately. Their motivation: full-matrix orthogonalization treats all heads as a single coupled block, so heads with larger gradient/momentum scales dominate the shared update direction while smaller-scale heads receive insufficiently normalized updates. Per-head orthogonalization equalizes update scale across heads, improves stability at larger scales, and is cheaper (Newton–Schulz on tall per-head blocks costs less than on the full projection).
  • Zhipu GLM-5 "Muon Split" (arXiv:2602.15763): the same idea applied to MLA up-projection matrices (split by attention head, orthogonalize each head independently), closing the performance gap between MLA and GQA under Muon.

Today in DeepSpeed, muon_update() (deepspeed/runtime/zero/muon/original_muon.py) applies Newton–Schulz to the full 2D parameter matrix — i.e., the "coupled block" behavior both papers describe.

Proposal

  1. Optimizer kernel: add a per-head mode to muon_update — reshape the update/momentum to [..., num_heads, head_dim] and run the existing bmm-based Newton–Schulz on the batched per-head blocks (the NS implementations already support batched inputs, so this is mostly a view/reshape branch). Apply the existing max(1, m/n)**0.5 scaling per head block.
  2. Parameter metadata pipeline: tag attention projection parameters with head structure (num_heads, head_dim, fused-vs-split QKV layout), similar to the existing use_muon tagging. The AutoTP model scan / presets already identify attention head structure, so this metadata can be sourced from there.
  3. ZeRO integration: the engine call site (post-unflatten, per-parameter muon_update) stays unchanged — just pass the head shape so the update can be viewed per-head and reshaped back before being applied to the flat buffer.
  4. Config surface: opt-in switch, e.g. optimizer.params.per_head_muon: true (or per param-group), following the precedent of separate muon_lr / adam_lr.
  5. Tests: unit tests for per-head vs full-matrix orthogonalization equivalence on synthetic attention-shaped params, plus convergence sanity on a small model.

References

Happy to help implement this.

Lenguaje dominante
Python
Estrellas
43.3k
Forks
5.1k
Merge medio
4 d 17 h
PR fusionados (30 d)
63

Preparar el entorno

Primeros pasos

  1. Lee el issue completo y luego la guía de contribución del proyecto.
  2. Comenta en el issue que vas a ocuparte — evita que dos personas hagan lo mismo.
  3. Haz un fork del repositorio y trabaja en una rama.
  4. Abre un pull request que haga referencia al número del issue.

Más de deepspeedai/DeepSpeed

Todos los issues de deepspeedai/DeepSpeed

Issues similares

Más issues de Python

Recibe los nuevos issues en tu correo

Un resumen breve de issues de GitHub para principiantes.