Skip to content

Entropy loss tweaks #465

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

Jump to bottom
Open
jlamypoirier wants to merge 4 commits into
base: main
Choose a base branch
from
Open

Entropy loss tweaks #465

jlamypoirier wants to merge 4 commits into from

Conversation

@jlamypoirier
Copy link
Collaborator

@jlamypoirier jlamypoirier commented Jan 29, 2026 •
edited
Loading

✨ Description

Content extracted from #428

Fix and improve LM losses

  • Tweak entropy losses to improve readability
  • Extract reusable fused_predicted_logits_from_labels from _fused_cross_entropy_base_from_labels
  • Support arbitrary tensor dimensions in entropy losses.
  • Add TP support for Z loss
  • Fix broken entropy loss tests for loss masking.
  • Generalize and simplify entropy loss tests with support for z loss, improve distributed coverage, add fp16 and bf16 tests.

Misc including:

  • Fix slow tokenization in tests
  • Fix some c10d imports
  • Automate cpu support for safe_barrier, etc.

@jlamypoirier jlamypoirier marked this pull request as ready for review January 30, 2026 04:00
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Reviewers

@oleksost oleksost Awaiting requested review from oleksost

@tscholak tscholak Awaiting requested review from tscholak

At least 1 approving review is required to merge this pull request.

Assignees

No one assigned

Labels

None yet

Projects

None yet

Milestone

No milestone

Development

Successfully merging this pull request may close these issues.

2 participants

@jlamypoirier