gather_cpu_state_dict¶
- torchtune.training.gather_cpu_state_dict(sharded_sd: Dict[str, DTensor], is_rank_zero: bool, device: Optional[device] = None) Dict[str, Any][源代码]¶
将分片的 state dict 转换为 CPU 上的完整 state dict。仅在 rank0 上返回非空结果，以避免 CPU 内存峰值。
- 参数:
sharded_sd (Dict[str, DTensor]) – DTensors 的分片 state dict
is_rank_zero (bool) – 标志，用于检查进程是否在 rank 0 上
device (Optional[torch.device]) – 用于分片张量的设备。默认值：None
- 返回:
CPU 上的 State dict
- 返回类型:
Dict[str, Any]
