Add a memory-efficient chunked LM-head log-probability utility to trl/trainer/utils.py and export it from…
Add a memory-efficient chunked LM-head log-probability utility to trl/trainer/utils.py and export it from…: a task in HF ML Tasksmith (Harbor dataset). ChunkedLogProbFunction — a torch.autograd.Function that computes per-token log-probabilities and entropy without ever materializing the full N, V…
The task
**`_ChunkedLogProbFunction`** — a `torch.autograd.Function` that computes per-token log-probabilities and entropy without ever materializing the full `[N, V]` logits matrix. It receives: - `last_hidden` `[N, H]` — flattened hidden states - `weight` `[V, H]` — LM head weight - `targets` `[N]` — token indices -…
Part of FineEnvs/HF_ML_Tasksmith.