Skip to content
New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

attribute speed and memory optimizations #245

Merged
merged 7 commits into from
Jan 17, 2024
Merged

attribute speed and memory optimizations #245

merged 7 commits into from
Jan 17, 2024

Conversation

gsarti
Copy link
Member

@gsarti gsarti commented Jan 10, 2024

Description

Small speed and memory optimizations to the main attribute loop, addressing issues #243 #240

  • Move tensors to CPU right away in the forward pass to avoid OOM when cloning
  • Fix remap_from_filtered behavior on sequence_scores tensors.
  • Use torch-native padding when converting lists of FeatureAttributionStepOutput to FeatureAttributionSequenceOutput in get_sequences_from_batched_steps.
  • Bump ruff version + update deps

@gsarti
Copy link
Member Author

gsarti commented Jan 10, 2024

Description of the problem of optimizing the get_sequences_from_batched_steps function:

You have a list of Pytorch tensors representing the results of computations performed at every step of a batched generation process involving one or more sequences. Every tensor in the list has a shape of [batch_size, sequence_length, ...], corresponding to scores for all previously generated steps for all sequences. Importantly, the sequences are not guaranteed to have the same length, so a decrease in batch_size between two consecutive tensors would denote the end of one or more sequences. Write a Python function to convert the list described above to a list of tensors of shape [max_sequence_length, num_sequence_steps, ...] representing the individual sequences, where max_sequence_length is the maximal length of the specific sequence, and num_sequence_steps is the number of steps for which scores are produced for the sequence.

@gsarti gsarti marked this pull request as ready for review January 17, 2024 12:15
@gsarti gsarti merged commit f434192 into main Jan 17, 2024
3 checks passed
@gsarti gsarti deleted the speed-mem-opt branch January 17, 2024 12:15
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment
Labels
enhancement New feature or request
Projects
None yet
1 participant