[tx] move logits computation LogitsProcessorMixin and fix prompt_logprobs

[raulchen]

Motivation

Different scenarios have different requirements for logits computation:

• Training: Only need logprobs of target tokens, not full [B, T, V] logits
• Prefill: Only need last position logits for sampling, plus optional prompt logprobs
• Decode: Need full [B, 1, V] logits for sampling transforms (temperature, top_k, top_p)

By returning hidden_states and letting callers decide what to compute, we avoid unnecessary materialization of
full vocabulary logits. This also enables future chunked logprobs computation for memory efficiency.

Changes

• Add LogitsProcessorMixin in tx/utils/logits_processor.py with:
  • Abstract get_lm_head() method
  • compute_logits(), compute_logprobs(), logits_to_logprobs()
• Remove logits from CausalLMOutput - callers now explicitly call compute_logits()
• Remove skip_prompt_logits parameter from model forward
• Add validation: train_unembed=True incompatible with tie_word_embeddings=True

Code consolidation

Logits/logprobs computation was scattered across multiple locations:

• tx/layers/logits_processor.py (deleted)
• tx/utils/generator.py (compute_prompt_logprobs() helper, manual log_softmax)
• tx/tinker/backends/jax.py (manual logprobs from logits)

Now unified in LogitsProcessorMixin with a single implementation.

Bug fix

The prompt_logprobs computation was not passing adapter_indices to compute_logprobs, causing incorrect results
with LoRA adapters. Added test coverage.

[gemini-code-assist]

Code Review

This is a well-executed refactoring that improves memory efficiency and code structure by separating the model's forward pass from the logits computation. The introduction of the CausalLMBase class effectively centralizes logic for computing logits and log probabilities, leading to cleaner and more maintainable code. The changes have been consistently applied across the different model implementations, the generator, and the training backend. I've identified a high-severity issue concerning LoRA with tied embeddings where adapter weights are not being applied to the final output projection. Additionally, there's a minor type-safety concern in one of the test files. Apart from these points, this is a solid pull request with clear benefits.

[gemini-code-assist]

There appears to be a bug here when using tied embeddings with LoRA adapters on the embed_tokens layer. The self.model.embed_tokens.T property returns a lambda function for the lm_head that computes hidden_states @ self.embedding.value.T. This uses the base embedding weights but ignores the LoRA adapter weights, as indicated by the TODO in LoRAEmbed.T. Consequently, the LoRA adaptation is not applied during the final projection to logits, leading to incorrect outputs. This should be addressed to ensure correctness for LoRA fine-tuning with tied embeddings.

[raulchen]

this is an existing bug / TODO.
added an assertion to prevent enabling tie_word_embedding and train_unembed together.

[raulchen]

/gemini review

[gemini-code-assist]

Code Review

This is an excellent pull request that significantly refactors the logits computation logic for better flexibility and memory efficiency. The introduction of the CausalLMBase class provides a clean and consolidated API for computing logits and logprobs, effectively removing duplication and simplifying the model's forward pass by eliminating the LogitsProcessor and the skip_prompt_logits parameter.

The changes are consistently applied across all models, tests, and utility functions. The generator logic is now more explicit about when and what kind of logits are computed, which is a great improvement for both prefill and decode steps. The adaptation in the JAX backend to use compute_logprobs within the gradient-checkpointed function is a smart move for memory optimization during training.

The tests have been thoughtfully updated to validate the new implementation, and the new check in init_lora_adapter to prevent misconfiguration with tied embeddings is a valuable addition. Overall, this is a high-quality, well-executed refactoring that improves the codebase's structure and performance characteristics.

[raulchen]

/gemini review

[gemini-code-assist]

Code Review

This pull request introduces a significant and well-executed refactoring of the logits and log-probabilities computation. By centralizing this logic into a LogitsProcessorMixin, the changes improve code clarity, remove redundancy, and fix a critical bug related to prompt_logprobs with LoRA adapters. The addition of a new test case to cover this fix and a validation check for incompatible configurations are also valuable improvements. My review includes one suggestion to optimize the prefill step to avoid redundant computations when prompt log-probabilities are requested.