Hacktoberfest 2026: le issue che i maintainer hanno segnato per ottobre, aperte e adatte ai principianti. Sfoglia le issue Hacktoberfest

Monolithic fused head-loss kernel for combinable losses + metrics

Aperta
  • Qualcosa non va?
#507 0 commenti 0 reazioni 0 assegnatari Vedi su GitHub

Nessuno ha ancora preso questa issue.

Valutazione

Difficoltà
5/5
Tempo stimato
Più di una settimana
Idoneità per principianti
12/100
Tipo di issue
Refactoring
Chiarezza
Da chiarire
Stato di attività
Tranquilla
Stack tecnologico
pytorch
Ambito
machine-learning, performance

Direzione di ricerca

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.

Scritto dal modello di indicizzazione a partire dal testo della issue.

Descrizione

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.

Lingua principale
Python
Stelle
337
Fork
52
Nessuna PR unita negli ultimi 30g

Preparare l'ambiente

Come iniziare

  1. Leggi tutta la issue e poi la guida ai contributi del progetto.
  2. Commenta sulla issue per dire che te ne occupi tu — evita che due persone facciano lo stesso lavoro.
  3. Fai un fork del repository e lavora su un branch.
  4. Apri una pull request che faccia riferimento al numero della issue.

Altre issue di ServiceNow/Fast-LLM

Tutte le issue di ServiceNow/Fast-LLM

Issue simili

Altre issue su Python

Ricevi le nuove issue nella tua casella

Un breve riepilogo di issue GitHub adatte ai principianti.