(new_model, module)
| 93 | def mp_merge_model_rank0(model, model_full): |
| 94 | assert get_model_parallel_world_size() == torch.distributed.get_world_size(), "Merging model is only supported for model_parallel_size == world_size!" |
| 95 | def iter_merge(new_model, module): |
| 96 | for (new_name, sub_new_model), (name, sub_module) in zip(new_model.named_children(), module.named_children()): |
| 97 | if isinstance(sub_module, (ColumnParallelLinear, RowParallelLinear, VocabParallelEmbedding)): |
| 98 | new_weights, new_biases = sub_module.partition() |
| 99 | new_weights = [x.to(sub_new_model.weight.device) for x in new_weights] |
| 100 | torch.distributed.gather(sub_new_model.weight.data, gather_list=new_weights, dst=0) |
| 101 | if new_biases: |
| 102 | new_biases = [x.to(sub_new_model.weight.device) for x in new_biases] |
| 103 | torch.distributed.gather(sub_new_model.bias.data, gather_list=new_biases, dst=0) |
| 104 | sub_module.merge([torch.clone(x.cpu()).detach() for x in new_weights], [torch.clone(x.cpu()).detach() for x in new_biases]) |
| 105 | del new_weights |
| 106 | if new_biases: |
| 107 | del new_biases |
| 108 | else: |
| 109 | for (nn, np), (n, p) in zip(sub_new_model.named_parameters(recurse=False), sub_module.named_parameters(recurse=False)): |
| 110 | p.data.copy_(torch.clone(np.data.cpu()).detach()) |
| 111 | iter_merge(sub_new_model, sub_module) |
| 112 | iter_merge(model, model_full) |
| 113 | |
| 114 | def mp_merge_model_send(model): |
no test coverage detected