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