🐛 Describe the bug
At line 160 of colossalai/zero/gemini/chunk/search_utils.py, inside method search_chunk_configuration, size_list is initialized using numel() of params. Then it's taken as the value of size_dict at line 168.
However, at line 175, the elements of size_list is compared with max_size, a variable with byte as its unit.
So I think at line 160, a factor of element_size should be multiplied when initializing size_list.
Environment
No response
🐛 Describe the bug
At line 160 of
colossalai/zero/gemini/chunk/search_utils.py, inside methodsearch_chunk_configuration,size_listis initialized using numel() of params. Then it's taken as the value ofsize_dictat line 168.However, at line 175, the elements of
size_listis compared withmax_size, a variable with byte as its unit.So I think at line 160, a factor of element_size should be multiplied when initializing
size_list.Environment
No response