(config)
| 12 | |
| 13 | |
| 14 | def create_empty_gptneox(config): |
| 15 | |
| 16 | import torch |
| 17 | import torch.nn as nn |
| 18 | |
| 19 | _reset_parameters_linear = nn.Linear.reset_parameters |
| 20 | def dummy(*args, **kargs): |
| 21 | pass |
| 22 | nn.Linear.reset_parameters = dummy |
| 23 | |
| 24 | # 1. disable init for faster initialization |
| 25 | # 2. avoid tie token embeddings with lm_head, as we train them separately. |
| 26 | with no_init_weights(_enable=True): |
| 27 | model = GPTNeoXForCausalLM(config).eval() |
| 28 | |
| 29 | nn.Linear.reset_parameters = _reset_parameters_linear |
| 30 | |
| 31 | return model |
| 32 | |
| 33 | def load_decentralized_checkpoint(model, checkpoint_path, n_stages=2, n_layer_per_stage=14): |
| 34 | input_path = checkpoint_path |
no outgoing calls
no test coverage detected