使用strategy.model_init_context()加载chatglm-6b模型提示mismatch

[zhangyuanscall]

加载代码如下所示，world_size 为4，使用单机4卡运行提示mismatch，transformer.word_embeddings.weight的维度为[130528,4096]，经过chunk之后成为[130528,1024]（即代码中x变量），但是p的维度为256（4096被切分为256），chatglm_model_path 为chatglm的huggingface权重

    chatglm_model_path = '/media/cfs/xxxx/LLMS/Pretrain_Models/chatglm_6b'
    state_dict = {}
    state_dict.update( torch.load(f"{chatglm_model_path}/pytorch_model-00001-of-00008.bin"))
    state_dict.update( torch.load(f"{chatglm_model_path}/pytorch_model-00002-of-00008.bin"))
    state_dict.update( torch.load(f"{chatglm_model_path}/pytorch_model-00003-of-00008.bin"))
    state_dict.update( torch.load(f"{chatglm_model_path}/pytorch_model-00004-of-00008.bin"))
    state_dict.update( torch.load(f"{chatglm_model_path}/pytorch_model-00005-of-00008.bin"))
    state_dict.update( torch.load(f"{chatglm_model_path}/pytorch_model-00006-of-00008.bin"))
    state_dict.update( torch.load(f"{chatglm_model_path}/pytorch_model-00007-of-00008.bin"))
    state_dict.update( torch.load(f"{chatglm_model_path}/pytorch_model-00008-of-00008.bin"))
    if is_rank_0():print(f"------------------> load state dict: {len(state_dict)} ")

    # configure model
    with strategy.model_init_context():
        if is_rank_0(): print('Warning: currently only bloom/chatglm is tested, gpt2,llama and opt are not tested')
        model = ChatGLMForConditionalGeneration(ChatGLMConfig.from_pretrained(chatglm_model_path))
        model.gradient_checkpointing_enable()
        world_size = dist.get_world_size()
        for n, p in model.named_parameters():
            x = state_dict[n]
            x = x.chunk(world_size, dim=-1)
            x = x[dist.get_rank()]
            print(f"----------------------------->n:{n},  x :{ x.shape }, p: {p.shape}")
            p.data.copy_(x)
        if is_rank_0(): print('---------------------Warning: load suceess')

报错日志如下：