Fix up query/key/value matrix ordering if checkpoint version is smaller than 2.0
(model, checkpoint_version)
| 247 | |
| 248 | |
| 249 | def fix_query_key_value_ordering(model, checkpoint_version): |
| 250 | """Fix up query/key/value matrix ordering if checkpoint |
| 251 | version is smaller than 2.0 |
| 252 | """ |
| 253 | if checkpoint_version < 2.0: |
| 254 | if isinstance(model, list): |
| 255 | assert len(model) == 1 |
| 256 | model = model[0] |
| 257 | for name, param in model.named_parameters(): |
| 258 | if name.endswith((".query_key_value.weight", ".query_key_value.bias")): |
| 259 | if checkpoint_version == 0: |
| 260 | fixed_param = _transpose_first_dim(param.data, 3, True, model) |
| 261 | elif checkpoint_version == 1.0: |
| 262 | fixed_param = _transpose_first_dim(param.data, 3, False, model) |
| 263 | else: |
| 264 | print_rank_0(f"Invalid checkpoint version {checkpoint_version}.") |
| 265 | sys.exit() |
| 266 | param.data.copy_(fixed_param) |
| 267 | if name.endswith((".key_value.weight", ".key_value.bias")): |
| 268 | if checkpoint_version == 0: |
| 269 | fixed_param = _transpose_first_dim(param.data, 2, True, model) |
| 270 | elif checkpoint_version == 1.0: |
| 271 | fixed_param = _transpose_first_dim(param.data, 2, False, model) |
| 272 | else: |
| 273 | print_rank_0(f"Invalid checkpoint version {checkpoint_version}.") |
| 274 | sys.exit() |
| 275 | param.data.copy_(fixed_param) |
| 276 | print_rank_0( |
| 277 | " successfully fixed query-key-values ordering for" |
| 278 | " checkpoint version {}".format(checkpoint_version) |
| 279 | ) |
| 280 | |
| 281 | |
| 282 | def load_deepspeed_state(model): |
no test coverage detected