[BUG]: Training process blocks on `gc.collect`

[HAOCHENYE]

🐛 Describe the bug

My training process was blocked when saving an optimizer state by Booster.save_optimizer. The program was blocked here:

ColossalAI/colossalai/zero/gemini/gemini_optimizer.py

Line 474 in 089c365

gc.collect()

Although I've no idea why the program was blocked at gc.collect, I believe that using gc.collect here is to free the storage as fast as possible. So, why don't we use the free_storage defined here:

ColossalAI/colossalai/zero/gemini/chunk/chunk.py

Line 43 in 089c365

def free_storage(tensor: torch.Tensor) -> None:

)

Replacing gc.collect with free_storage solves my problem perfectly.

If replacing gc.collect with free_storage is reasonable, I'd be happy to create a PR to fix it.

Environment

No response

[flybird11111]

free_storage is a specific method in PyTorch for releasing tensor memory, while gc.collect is part of Python's garbage collection mechanism for releasing memory of unused objects. Thanks for tuning in.

[HAOCHENYE]

I understand that gc.collect is used to manually trigger garbage collection in Python. However, as I mentioned above, I believe the purpose of this line of code is to free the tensor memory (according to the type hint, state_shard[0] should be a Tensor or NoneObject), so would using free_storage be more appropriate? Another consideration is that my program hangs during the gc.collect phase, but replacing it with free_storage resolves the issue.

[Fridge003]

Hi, I'm the author of this piece of code. I have created a PR #4396 to fix this issue according to your suggestion. Thank you for your contribution!