Hacktoberfest 2026: những issue maintainer đã đánh dấu cho tháng Mười, đang mở và phù hợp người mới. Xem issue Hacktoberfest

Monolithic fused head-loss kernel for combinable losses + metrics

Đang mở
  • Có gì không ổn?
#507 0 bình luận 0 reaction 0 người được giao Xem trên GitHub

Chưa có ai nhận issue này.

Đánh giá

Độ khó
5/5
Thời gian dự kiến
Hơn một tuần
Mức phù hợp với người mới
12/100
Loại issue
Tái cấu trúc
Độ rõ ràng
Cần làm rõ
Mức độ hoạt động
Ít trao đổi
Công nghệ
pytorch
Lĩnh vực
machine-learning, performance

Hướng nghiên cứu

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.

Do mô hình lập chỉ mục viết ra từ nội dung của issue.

Mô tả

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.

Ngôn ngữ chính
Python
Star
337
Fork
52
Không có pull request nào được merge trong 30 ngày

Chuẩn bị môi trường

Bắt đầu từ đâu

  1. Đọc hết issue, rồi đọc hướng dẫn đóng góp của dự án.
  2. Bình luận trên issue rằng bạn sẽ nhận — tránh hai người làm cùng một việc.
  3. Fork repository và làm thay đổi trên một nhánh.
  4. Mở pull request có tham chiếu số hiệu của issue.

Issue khác của ServiceNow/Fast-LLM

Tất cả issue của ServiceNow/Fast-LLM

Issue tương tự

Thêm issue về Python

Nhận issue mới trong hộp thư của bạn

Bản tóm tắt ngắn những issue GitHub phù hợp với người mới.