[Gemini] fix the convert_to_torch_module bug

[feifeibear]

What's new

Previously, convert_to_torch_module dose not work. Because it can not find the correct place to store param.data. Now, I fix this bug.

You can test the feature with the following code.

torchrun --standalone --nproc_per_node=1 XXX.py

import colossalai
from colossalai.utils import get_current_device
from colossalai.nn.parallel import GeminiDDP
from transformers import BloomForCausalLM, AutoTokenizer
from colossalai.utils.model.colo_init_context import ColoInitContext, post_process_colo_init_ctx
from colossalai.tensor import ColoTensor, ProcessGroup
from colossalai.nn.parallel.utils import convert_to_torch_module


config = {
        "BATCH": 4,
        "gradient_accumulation_steps": 1,
        "clip_grad_norm": 1,
    }

colossalai.launch_from_torch(config=config)

pg = ProcessGroup()

with ColoInitContext():
    model = BloomForCausalLM.from_pretrained("/data2/users/lczht/bloom-560m")

param_num = len([p for n, p in model.named_parameters()])
print(f'param num {param_num}')
# for p in model.parameters():
#    print(p)

# after wrap with GeminiDDP model parameter is broken.
model = GeminiDDP(model,
                device=get_current_device(),
                placement_policy="cuda",
                pin_memory=True,
                search_range_mb=32)

# for n, p in model.torch_named_parameters():
#     print(n, p)

# for k, v in modelv2.state_dict().items():
#     print(k, v)

model = convert_to_torch_module(model)

cnt = 0
for name, param in model.named_parameters(recurse = True):
    if isinstance(param, ColoTensor):
        print('meet ColoTensor', name)
    cnt += 1
print(f'{cnt} params found')