Use runtime profiling to replace manual memory analyzers

[zhuohan123]

Fix #59.

Previously we used a manual memory profiler, which must be implemented separately for each model. This PR replaces it with a general memory profiler.

[zhuohan123]

@WoosukKwon This PR is ready for review.

[WoosukKwon]

Thanks @zhuohan123 for this PR. Please check my comments.

[WoosukKwon]

Additionally, I've found that the PR changes the output of the examples in simple_server.py. We should reset the random generator states after the profiling run.

[WoosukKwon]

@zhuohan123 I will start to merge the other PRs first, so that this PR does not block any.

[zhuohan123]

@WoosukKwon I fixed all the review comments and merged with the latest master. Please review again. Feel free to merge it.

One thing I noticed is this bug in the sampler. With this bug, I discover that these two sorts (link1, link2) take a lot of memory (4-5GB) when sampling for 2560 tokens at the same time. In this case, they are sorting a (2560, 51200) shape tensor. I'm not sure why these two sorts take that much memory and whether there are more memory efficient way to implement topp and topk.

[zhuohan123]

Additionally, I've found that the PR changes the output of the examples in simple_server.py. We should reset the random generator states after the profiling run.

I haven't added it in the current PR since it looks pretty redundant to add an additional "set_random_seed" call to all the workers after profiling. It will be a nested ray call and complicates the code. I think this will confuse future people when they read the code.

[WoosukKwon]

I haven't added it in the current PR since it looks pretty redundant to add an additional "set_random_seed" call to all the workers after profiling. It will be a nested ray call and complicates the code. I think this will confuse future people when they read the code.

What do you mean by "nested ray call"? I didn't actually get it. And I think the purpose of resetting the random state is to make sure everything that happens before starting serving should not affect the outputs. In this sense, I think we should reset the random states after the profiling run, while we don't have to reset the states before profiling run.

[WoosukKwon]

One thing I noticed is this bug in the sampler. With this bug, I discover that these two sorts (link1, link2) take a lot of memory (4-5GB) when sampling for 2560 tokens at the same time. In this case, they are sorting a (2560, 51200) shape tensor. I'm not sure why these two sorts take that much memory and whether there are more memory efficient way to implement topp and topk.

1. Thanks for finding the bug. As you did, -1 should be replaced with vocab_size.
2. On the memory consumption, it seems it's a known issue: Peak GPU-memory usage extremely huge when sorting with torch.sort pytorch/pytorch#77049. For now, I think we can take into account max_num_sequences (which is 256 by default) in calculating the peak memory usage. This will reduce the memory consumption of sorting by 10x.

[WoosukKwon]

LGTM. Thanks!