gather_cpu_state_dict

torchtune.training.gather_cpu_state_dict(model: FSDPModule, is_rank_zero: bool, device: Optional[device] = None, adapter_weights_only: bool = False) → Dict[str, Any]

将分片的 state dict 转换为 CPU 上的完整 state dict，仅在 rank0 上返回非空结果以避免 CPU 内存峰值。当前我们可以使用分布式 state dict API 处理不包含 NF4Tensor 的模型。否则，我们需要手动收集所有 NF4 张量，直到 NF4Tensor 子类支持 all-gather。TODO：在分布式 state dict API 中添加对 NF4Tensor 的支持

参数：

• model (FSDPModule) – 要为其生成 cpu_state_dict 完整限定名称的模型

• is_rank_zero (bool) – 检查进程是否在 rank 0 上的标志

• device (Optional[torch.device]) – 用于分片张量的设备。默认值：None

• adapter_weights_only (bool) – 检查是否仅返回可训练参数的标志。默认值：False

返回值：

CPU 上的 State dict

返回类型：

Dict[str, Any]