torch.distributed.all_reduce does not free memory

[WoosukKwon]

I've visualized the memory usage:

• llama 7B, TP=1

The activation memory is reused after every layer.

• llama-70B, TP=8

However, when using TP, the activation memory for all reduce is not reused

Originally posted by @WoosukKwon in #2031 (comment)