<think>好的，用户的问题是关于如何保存和加载使用PyTorch FSDP训练的模型和优化器。我需要先回忆一下FSDP的基本原理和保存加载的注意事项。FSDP（Fully Sharded Data Parallel）是一种分布式训练技术，它将模型参数、梯度和优化器状态分片到各个GPU上，这样可以节省内存，但同时也让保存和加载变得复杂。 首先，保存模型的时候，普通的方法比如`torch.save`可能无法正确保存分片的参数，因为每个GPU只保存了一部分。需要使用FSDP提供的专门方法，比如`state_dict`的获取和保存。记得在保存模型状态字典时，应该只在rank 0上保存，避免重复。同时，可能需要使用`dist.barrier()`来同步进程，防止文件写入冲突。 然后是优化器的保存。优化器的状态同样被分片了，所以也需要使用FSDP特定的方法。可能要用到`optim.state_dict()`，同样需要收集所有分片的状态，并且只在主进程保存。加载的时候，要确保在初始化模型和优化器之后，正确地加载这些状态字典，可能需要使用`scatter`方法将状态分布到各个GPU。 另外，用户提供的引用中有提到合并PEFT模型和基础模型的部分，比如使用`AutoPeftModelForCausalLM`的`merge_and_unload`方法。如果用户还使用了LoRA等参数高效微调技术，保存时可能需要先合并模型再保存，或者保存适配器权重，这部分需要额外注意。 用户可能还关心如何在不中断训练的情况下进行断点续训，所以需要保存优化器状态、学习率调度器等。需要确认FSDP是否支持这些部分的保存和加载，以及具体的API调用方式。 最后，要提醒用户注意内存和存储的问题，比如保存分片模型可能需要更多的CPU内存，尤其是在合并模型时，引用[2]提到可能需要超过192GB的CPU内存。此外，保存的路径需要是可访问的，可能需要使用分布式文件系统或者只在主进程保存。 总结步骤应该是：1. 保存模型时使用FSDP的state_dict，并在rank 0处理；2. 保存优化器状态同样处理；3. 加载时正确加载状态字典并分布到各GPU；4. 处理可能的PEFT合并问题；5. 注意内存和进程同步问题。</think>在PyTorch中使用FSDP（Fully Sharded Data Parallel）保存和加载模型与优化器需要注意分片参数的特性，以下是具体实现方法： ### 一、保存FSDP模型 1. **获取模型状态字典** 使用`FSDP.state_dict()`接口获取分片参数的状态字典，需配合`state_dict_type`参数控制保存粒度： ```python from torch.distributed.fsdp import FullyShardedDataParallel as FSDP # 设置状态字典类型（FULL_STATE_DICT会聚合所有分片） with FSDP.state_dict_type(model, StateDictType.FULL_STATE_DICT): model_state = model.state_dict() ``` 2. **仅主进程保存文件** 通过`dist.get_rank() == 0`判断主进程，避免多进程重复写入： ```python if dist.get_rank() == 0: torch.save(model_state, "model.pth") dist.barrier() # 同步所有进程 ``` ### 二、保存FSDP优化器 1. **获取优化器状态字典** 优化器状态也需特殊处理分片特性： ```python optim_state = FSDP.optim_state_dict(model, optimizer) ``` 2. **分片状态保存** 直接保存可能包含分布式信息，建议使用专用API： ```python if dist.get_rank() == 0: torch.save(optim_state, "optimizer.pth") ``` ### 三、加载FSDP模型 1. **初始化空模型** 先构建未经分片的原始模型结构： ```python model = MyModel().to(device) ``` 2. **加载并转换状态字典** 使用`FSDP.shard_state_dict`处理分片映射： ```python loaded_model_state = torch.load("model.pth") sharded_state = FSDP.shard_state_dict(loaded_model_state) model.load_state_dict(sharded_state) ``` ### 四、加载FSDP优化器 1. **重建优化器** 先创建与原配置相同的优化器： ```python optimizer = torch.optim.AdamW(model.parameters()) ``` 2. **分片加载优化器状态** 使用`FSDP.optim_state_dict_to_load`转换状态： ```python loaded_optim_state = torch.load("optimizer.pth") optim_state = FSDP.optim_state_dict_to_load(model, optimizer, loaded_optim_state) optimizer.load_state_dict(optim_state) ``` ### 五、特殊场景处理 - **混合精度训练**：保存时需保留`dtype`信息 - **多节点训练**：确保文件系统支持并行访问（如NFS） - **断点续训**：需同时保存epoch计数等附加信息 > 重要说明：若使用PEFT（如LoRA），需先合并参数再保存完整模型[^2]，否则应单独保存适配器权重： > ```python > merged_model = model.merge_and_unload() > merged_model.save_pretrained("output_dir") # 如引用[2]所示 > ```