Add logits_to_keep parameter to BertForSequenceClassification

[Kraga922]

Summary

This PR adds support for the logits_to_keep parameter to BertForSequenceClassification

Motivation

The logits_to_keep parameter enables memory-efficient inference by computing logits only for the last N token positions. This optimization is particularly useful for:

• Generation tasks where only recent tokens are needed
• Memory-constrained environments
• Processing long sequences efficiently

Implementation Details

• Added logits_to_keep parameter to the forward() method signature
• Implemented token slicing logic that activates when parameter is specified
• Updated docstring with comprehensive parameter documentation
• Maintains full backward compatibility (parameter defaults to None)

Code Changes

The implementation adds approximately 10 lines of code:

1. New optional parameter in function signature
2. Conditional slicing logic for sequence outputs
3. Documentation in docstring

Testing

Comprehensive local testing performed with bert-base-uncased:

• Single input and batch processing verified
• Multiple logits_to_keep values tested (1, 2, 3, 5, 10, None)
• Edge cases handled (0, negative values, values > sequence length)
• Training compatibility confirmed (gradient flow works)
• Output structure validation passed
• Small performance improvement observed (~0.6% faster)

Backward Compatibility

• Parameter is optional and defaults to None
• When None, behavior is identical to current implementation
• No breaking changes to existing code

Related Issues

Inspired by #40984

[github-actions]

[For maintainers] Suggested jobs to run (before merge)

run-slow: bert

[Rocketknight1]

Hi @Kraga922, I'm going to close this in favour of #41335. Although we appreciate the PR, it's easier to review a big block of similar changes together!