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

Monolithic fused head-loss kernel for combinable losses + metrics

Abierto
  • ¿Algo no está bien?
#507 0 comentarios 0 reacciones 0 asignados Ver en GitHub

Nadie ha tomado este issue todavía.

Evaluación

Dificultad
5/5
Tiempo estimado
Más de una semana
Aptitud para principiantes
12/100
Tipo de issue
Refactorización
Claridad
Necesita aclaración
Estado de actividad
Tranquilo
Stack tecnológico
pytorch
Área
machine-learning, performance

Línea de trabajo

The redundant softmax passes live in fused_grpo_loss_forward_backward, compute_policy_gradient_metrics, the fused_cross_entropy_base_from* helpers and fused_softmax_base, with the entropy chunking from PR #494. Start by reading those functions to map where sum_exp_logits and the tensor-parallel all-reduces are computed. The issue defers the design, so done needs an agreed design with the maintainers before any implementation.

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

Descripción

Problem

Each language-model head loss/metric runs its own softmax pass over the full vocab. When losses are combined, the work is duplicated:

  • cross-entropy + z-loss — both compute sum_exp_logits over the same logits.
  • cross-entropy + distillation — both compute the student softmax (and additionally the teacher softmax for distillation).
  • GRPO + extra metrics (PR #494) — fused_grpo_loss_forward_backward computes logits_norm, exp_logits, sum_exp_logits, and predicted_logits; compute_policy_gradient_metrics then recomputes all of them on the same logits. With compute_entropy_metric=True, a third softmax pass runs on top.

Each pass also issues its own tensor-parallel all-reduces on logits_max / sum_exp_logits, multiplying the comm cost.

@torch.compile does not fuse across separate decorated functions, so the redundant work is real both in compute and memory.

Proposed direction

A single "monolithic" head-loss kernel (torch.compile and/or triton) that:

  • Takes a config / flag-set describing which losses and metrics to emit (CE, z-loss, distillation, GRPO clipped objective, GRPO ratio/KL/clamp/advantage stats, entropy, ...).
  • Runs softmax (and TP all-reduce) once over the logits.
  • Emits all requested scalars and the combined gradient in one pass.

This subsumes fused_grpo_loss_forward_backward, _fused_cross_entropy_base_from_*, fused_softmax_base, compute_policy_gradient_metrics, and the entropy chunking in PR #494.

Out of scope

  • Implementation details (triton vs torch.compile, kernel layout) — left for the design phase.
  • The current PRs that hit this redundancy (notably #494) are not blocked on this issue and should land with their existing structure; this is the longer-term consolidation.

Motivation

Surfaced during review of #494, but the underlying redundancy predates it.

Lenguaje dominante
Python
Estrellas
337
Forks
52
Sin PR fusionados en 30 d

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 ServiceNow/Fast-LLM

Todos los issues de ServiceNow/Fast-LLM

Issues similares

Más issues de Python

Recibe los nuevos issues en tu correo

Un resumen breve de issues de GitHub para principiantes.