| 79 | model.load_state_dict(sharded_sd, strict=False, assign=True) |
| 80 | |
| 81 | def load_optim(self, model: FSDPModule, opt: torch.optim.Optimizer): |
| 82 | last_optim_checkpoint = ( |
| 83 | f"{self.folder}/{'dcp_api' if self.dcp_api else 'dtensor_api'}" |
| 84 | f"/{self.last_training_time}/{OPTIM_CHECKPOINT}" |
| 85 | ) |
| 86 | full_sd = torch.load( |
| 87 | last_optim_checkpoint, mmap=True, weights_only=True, map_location="cpu" |
| 88 | ) |
| 89 | if self.dcp_api: |
| 90 | set_optimizer_state_dict( |
| 91 | model=model, |
| 92 | optimizers=opt, |
| 93 | optim_state_dict=full_sd, |
| 94 | options=StateDictOptions( |
| 95 | full_state_dict=True, |
| 96 | broadcast_from_rank0=True, |
| 97 | ), |
| 98 | ) |
| 99 | return |
| 100 | _init_optim_state(opt) |
| 101 | param_groups = opt.state_dict()["param_groups"] |
| 102 | state = opt.state_dict()["state"] |
| 103 | |
| 104 | full_param_groups = full_sd["param_groups"] |
| 105 | full_state = full_sd["state"] |
| 106 | |
| 107 | for param_group, full_param_group in zip(param_groups, full_param_groups): |
| 108 | for key, value in full_param_group.items(): |
| 109 | if key == PARAMS: |
| 110 | continue |
| 111 | param_group[key] = value |
| 112 | for pid, full_pid in zip(param_group[PARAMS], full_param_group[PARAMS]): |
| 113 | if pid not in state: |
| 114 | continue |
| 115 | param_state = state[pid] |
| 116 | full_param_state = full_state[full_pid] |
| 117 | for attr, full_tensor in full_param_state.items(): |
| 118 | sharded_tensor = param_state[attr] |
| 119 | if isinstance(sharded_tensor, DTensor): |
| 120 | # exp_avg is DTensor |
| 121 | param_state[attr] = distribute_tensor( |
| 122 | full_tensor, |
| 123 | sharded_tensor.device_mesh, |
| 124 | sharded_tensor.placements, |
| 125 | ) |
| 126 | else: |
| 127 | # step is plain tensor |
| 128 | param_state[attr] = full_tensor |
| 129 | opt.load_state_dict( |
| 130 | { |
| 131 | "param_groups": param_groups, |
| 132 | "state": state, |
| 133 | } |
| 134 | ) |
| 135 | |
| 136 | def _get_full_model_state_dict(self, model: FSDPModule): |
| 137 | if self.dcp_api: |