(module, state_dict, prefix, mp_group=None)
| 157 | module._buffers[name].data.copy_(state_dict[prefix + name]) |
| 158 | |
| 159 | def load(module, state_dict, prefix, mp_group=None): |
| 160 | mp_replace = ReplaceWithTensorSlicing(mp_group=mp_group) |
| 161 | if hasattr(module, 'weight'): |
| 162 | if module.weight.data.is_meta: |
| 163 | # meta tensor cannot be casted or copied to, so we need to replace it with a normal tensor here |
| 164 | module.weight = torch.nn.parameter.Parameter(data=torch.empty_like(module.weight.data, device="cpu"), |
| 165 | requires_grad=module.weight.data.requires_grad) |
| 166 | if 'query_key_value' in prefix: |
| 167 | module.weight = mp_replace.strided_copy(module.weight.data, |
| 168 | state_dict[prefix + 'weight'], |
| 169 | num_splits=3) |
| 170 | else: |
| 171 | module.weight = mp_replace.copy(module.weight.data, state_dict[prefix + 'weight']) |
| 172 | else: |
| 173 | if hasattr(module, 'norm') and hasattr(module.norm, 'weight'): |
| 174 | if module.norm.weight.data.is_meta: |
| 175 | # meta tensor cannot be casted or copied to, so we need to replace it with a normal tensor here |
| 176 | module.norm.weight = torch.nn.parameter.Parameter( |
| 177 | data=torch.empty_like(module.norm.weight.data, device="cpu"), |
| 178 | requires_grad=module.norm.weight.data.requires_grad) |
| 179 | module.norm.weight = mp_replace.copy(module.norm.weight.data, state_dict[prefix + 'weight']) |
| 180 | |
| 181 | if prefix + 'bias' in state_dict.keys(): |
| 182 | if hasattr(module, 'bias'): |
| 183 | if module.bias.data.is_meta: |
| 184 | # meta tensor cannot be casted or copied to, so we need to replace it with a normal tensor here |
| 185 | module.bias = torch.nn.parameter.Parameter(data=torch.empty_like(module.bias.data, device="cpu"), |
| 186 | requires_grad=module.bias.data.requires_grad) |
| 187 | module.bias = mp_replace.copy(module.bias, state_dict[prefix + 'bias']) |
| 188 | else: |
| 189 | if hasattr(module, 'norm') and hasattr(module.norm, 'bias'): |
| 190 | if module.norm.bias.data.is_meta: |
| 191 | # meta tensor cannot be casted or copied to, so we need to replace it with a normal tensor here |
| 192 | module.norm.bias = torch.nn.parameter.Parameter( |
| 193 | data=torch.empty_like(module.norm.bias.data, device="cpu"), |
| 194 | requires_grad=module.norm.bias.data.requires_grad) |
| 195 | module.norm.bias = mp_replace.copy(module.norm.bias, state_dict[prefix + 'bias']) |
| 196 | |
| 197 | |
| 198 | class AutoTP(): |
no test coverage detected