monkeypatch.trainer.utils
monkeypatch.trainer.utils
Functions
| Name | Description |
|---|---|
| entropy_from_logits | Triton-fused entropy (online single-pass). Handles non-contiguous tensors without copying. |
| selective_log_softmax | Fused selective_log_softmax with Triton forward+backward kernels. |
| selective_log_softmax_original | Original selective_log_softmax (reference/fallback). |
entropy_from_logits
monkeypatch.trainer.utils.entropy_from_logits(logits, chunk_size=128)Triton-fused entropy (online single-pass). Handles non-contiguous tensors without copying.
selective_log_softmax
monkeypatch.trainer.utils.selective_log_softmax(logits, index)Fused selective_log_softmax with Triton forward+backward kernels.
Equivalent to: torch.gather(logits.log_softmax(-1), dim=-1, index=index)
selective_log_softmax_original
monkeypatch.trainer.utils.selective_log_softmax_original(logits, index)Original selective_log_softmax (reference/fallback).