[megatron] Fused LM-head log-prob + entropy (avoid full [*, seq, vocab] logit materialization)#1765

Open

dyurk-lila wants to merge 7 commits into

NovaSky-AI:mainfrom

dyurk-lila:feat/fused-linear-cross-entropy

Commits on Jun 8, 2026

feat(megatron): fused linear log-prob (hidden->logprob) to avoid materializing [*, seq, vocab] logits

dyurk-lila
and
claude
committed
fix(megatron): Triton fused-logprob OOV gradient + GPU-validated precision policy

dyurk-lila
and
claude
committed
fix(megatron): cast hidden to LM-head weight dtype in Triton fused-logprob adapter

dyurk-lila
and
claude
committed
feat(megatron): fused entropy dual-output → wire SFT + ALL RL losses through the fused path

dyurk-lila
and
claude
committed
test(megatron): initialize mpu model-parallel state in the Triton entropy GPU test

dyurk-lila
and
claude
committed

Commits on Jun 9, 2026