This function makes partitions in-place for a model. It takes less memory when world size is small.
(model, new_model_parallel_size)
| 24 | raise NotImplementedError |
| 25 | |
| 26 | def mp_split_model(model, new_model_parallel_size): |
| 27 | """ |
| 28 | This function makes partitions in-place for a model. |
| 29 | It takes less memory when world size is small. |
| 30 | """ |
| 31 | from sat.model.transformer import SelfAttention, CrossAttention |
| 32 | |
| 33 | destroy_model_parallel() |
| 34 | initialize_model_parallel(new_model_parallel_size) |
| 35 | def iter_repartition(module): |
| 36 | for name, sub_module in module.named_children(): |
| 37 | if isinstance(sub_module, (ColumnParallelLinear, RowParallelLinear, VocabParallelEmbedding, |
| 38 | SelfAttention, CrossAttention)): |
| 39 | sub_module.repartition() |
| 40 | iter_repartition(sub_module) |
| 41 | iter_repartition(model) |
| 42 | |
| 43 | from .initialize import get_node_group, get_node_src_rank, get_node_world_size |
| 44 | from .initialize import get_model_parallel_group, get_model_parallel_src_rank, get_model_parallel_world_size |