Deepspeed Zero3下如何用Safetensors完整保存含冻结参数的微调模型？

我尝试通过仅训练lm_head层微调因果语言模型（Casual LM），使用Deepspeed Zero Stage3训练过程正常，但保存微调后的模型并重新加载时遇到问题。执行保存代码后仅保存了未冻结的lm_head部分（仅7MB，原模型超30GB），加载时出现尺寸不匹配错误。

微调代码 #

accelerator = Accelerator(log_with="tensorboard", project_dir=project_dir)
model: torch.nn.Module = AutoModelForCausalLM.from_pretrained("the path to LM model", trust_remote_code=True)
model.half()
model.train()

# 冻结所有参数，仅解冻lm_head
for param in model.parameters():
    param.requires_grad = False
for param in model.lm_head.parameters():
    param.requires_grad = True  

...

# 保存微调模型
if step == 5000 and accelerator.is_main_process:
     unwrapped_model: PreTrainedModel = accelerator.unwrap_model(model)
     save_fn = accelerator.save
     unwrapped_model.save_pretrained(
         "mycogagent",
         is_main_process=accelerator.is_main_process,
         save_function=save_fn,
     )

保存时警告 #

Removed shared tensor {a long list of parameter names in the original LM model except the parameter name of lm_head} while saving. This should be OK, but check by verifying that you don't receive any warning while reloading

加载错误 #

执行加载代码：

model = AutoModelForCausalLM.from_pretrained("mycogagent", trust_remote_code=True)

报错：

RuntimeError: Error(s) in loading state_dict for CogAgentForCausalLM:
        size mismatch for model.embed_tokens.weight: copying a param with shape torch.Size([0]) from checkpoint, the shape in current model is torch.Size([32000, 4096]).
        You may consider adding `ignore_mismatched_sizes=True` in the model `from_pretrained` method.

添加ignore_mismatched_sizes=True后，仍出现大量初始化警告并停滞：

Some weights of CogAgentForCausalLM were not initialized from the model checkpoint at mycogagent and are newly initialized:[a list of parameter names]
You should probably TRAIN this model on a down-stream task to be able to use it for predictions and inference.

Some weights of CogAgentForCausalLM were not initialized from the model checkpoint at mycogagent and are newly initialized because the shapes did not match:
- model.embed_tokens.weight: found shape torch.Size([0]) in the checkpoint and torch.Size([32000, 4096]) in the model instantiated
You should probably TRAIN this model on a down-stream task to be able to use it for predictions and inference.

问题根源定位 #

保存过程中state_dict看似正常，但transformers/pytorch_utils.py中的id_tensor_storage(tensor)函数无法获取张量正确指针，输出始终为(device(type='cuda', index=0), 0, 0)，而正常情况下unique_id应为张量内存地址，问题出在accelerate.unwrap函数对Deepspeed Zero3下冻结参数的处理上。

Deepspeed Zero Stage3会对未训练（冻结）的参数进行分片或不在主进程保留完整副本，导致accelerator.unwrap_model无法恢复完整参数，进而save_pretrained仅保存了可训练的lm_head。以下是两种可靠解决方法：

方法1：主进程加载原始模型并替换微调后的lm_head #

在主进程中加载原始预训练模型，将微调后的lm_head权重替换进去后保存完整模型：

if step == 5000 and accelerator.is_main_process:
    # 取出微调后的模型
    unwrapped_model = accelerator.unwrap_model(model)
    # 加载原始预训练模型（和微调时的模型路径一致）
    full_model = AutoModelForCausalLM.from_pretrained(
        "the path to LM model",
        trust_remote_code=True,
        torch_dtype=torch.float16
    )
    # 替换lm_head的权重
    full_model.lm_head.load_state_dict(unwrapped_model.lm_head.state_dict())
    # 保存完整模型
    full_model.save_pretrained("mycogagent")

方法2：强制收集所有参数到主进程后保存 #

利用Deepspeed的gather功能，将所有参数（包括冻结的）聚集到主进程，再生成完整state_dict保存：

if step == 5000 and accelerator.is_main_process:
    unwrapped_model = accelerator.unwrap_model(model)
    # 生成完整的state_dict，包含所有参数
    full_state_dict = {}
    for name, param in unwrapped_model.named_parameters():
        # 收集参数到主进程
        gathered_param = accelerator.gather(param).cpu()
        full_state_dict[name] = gathered_param
    # 加载原始模型并加载完整state_dict
    full_model = AutoModelForCausalLM.from_pretrained(
        "the path to LM model",
        trust_remote_code=True,
        torch_dtype=torch.float16
    )
    full_model.load_state_dict(full_state_dict, strict=True)
    full_model.save_pretrained("mycogagent")

关键注意事项 #

• 不要使用accelerator.save作为save_pretrained的save_function，Zero3下该方法仅处理可训练参数；
• 冻结参数在Zero3中不会被主进程完整持有，必须从原始模型加载或手动聚集；
• 保存时确保主进程持有完整的模型权重，避免仅保存部分参数。

内容的提问来源于stack exchange，提问作者Yaoming Xuan