Skip to content
This repository has been archived by the owner on Nov 3, 2023. It is now read-only.

Parallelize detensorizing context over batch #3729

Merged
merged 3 commits into from
Jun 18, 2021

Conversation

EricMichaelSmith
Copy link
Contributor

@EricMichaelSmith EricMichaelSmith commented Jun 16, 2021

Patch description
Currently, when generating, we convert the context Tensor to a list one example at a time, which is expensive; this PR parallelizes this operation over the entire batch.

Performance on 1000 generations (-mf zoo:blender/blender_90M/model -t blended_skill_talk -ne 1000), on an otherwise unoccupied devfair:

  • Original code, trial 1: median elapsed time 783 ms, mean 797 ms
  • Original code, trial 2: median 770 ms, mean 785 ms
  • This PR, trial 1: median 651 ms, mean 663 ms
  • This PR, trial 2: median 665 ms, mean 677 ms

Testing steps
CI checks (which seem to cover context blocking)

Copy link
Contributor

@stephenroller stephenroller left a comment

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

This one seems great to me... Haven't reviewed the other. We can talk more offline.

"""
if self.beam_context_block_ngram <= 0:
# We aren't context blocking, return empty tensor of the correct size
return torch.LongTensor([[]] * batch.batchsize)
Copy link
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Suggested change
return torch.LongTensor([[]] * batch.batchsize)
return torch.zeros(batch.batchsize, 0)

Copy link
Contributor Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Yeah, that's cleaner - just added that, but with dtype=torch.long

@EricMichaelSmith EricMichaelSmith merged commit e853e7c into master Jun 18, 2021
@EricMichaelSmith EricMichaelSmith deleted the context-blocking-speedups branch June 18, 2021 20:28
Sign up for free to subscribe to this conversation on GitHub. Already have an account? Sign in.
Projects
None yet
Development

Successfully merging this pull request may close these issues.

3 participants