|
| 1 | +# Copyright (c) Meta Platforms, Inc. and affiliates. |
| 2 | +# All rights reserved. |
| 3 | +# |
| 4 | +# This source code is licensed under the terms described in the LICENSE file in |
| 5 | +# the root directory of this source tree. |
| 6 | + |
| 7 | +import concurrent.futures |
| 8 | +import re |
| 9 | +from pathlib import Path |
| 10 | +from typing import Any, Dict, List, Optional, Union |
| 11 | + |
| 12 | +import numpy as np |
| 13 | +import torch |
| 14 | +from fairscale.nn.model_parallel.initialize import get_model_parallel_rank, get_model_parallel_world_size |
| 15 | + |
| 16 | + |
| 17 | +def map_mp_rank(old_mp_size: int, new_mp_size: int, new_mp_rank: int) -> List[int]: |
| 18 | + """Map a new MP rank to a list of old MP ranks given a change in MP size.""" |
| 19 | + if new_mp_size % old_mp_size == 0: |
| 20 | + # Read old MP shard and split it into smaller ones |
| 21 | + return [new_mp_rank * old_mp_size // new_mp_size] |
| 22 | + elif old_mp_size % new_mp_size == 0: |
| 23 | + # Merge old MP shards into a single one |
| 24 | + mp_factor = old_mp_size // new_mp_size |
| 25 | + return list(range(new_mp_rank * mp_factor, (new_mp_rank + 1) * mp_factor)) |
| 26 | + else: |
| 27 | + raise ValueError( |
| 28 | + f"Either old MP size or new MP size should be a multiple of the other: " |
| 29 | + f"{old_mp_size} % {new_mp_size} != 0 and {new_mp_size} % {old_mp_size} != 0" |
| 30 | + ) |
| 31 | + |
| 32 | + |
| 33 | +def maybe_reshard_state_dict( |
| 34 | + ckpt_paths: List[Path], |
| 35 | + n_kv_heads: int, |
| 36 | + moe_num_experts: Optional[int] = None, |
| 37 | + map_location: Union[str, torch.device] = "cpu", |
| 38 | + mmap: bool = True, |
| 39 | +) -> Dict[str, torch.Tensor]: |
| 40 | + if str(map_location) == "cpu": |
| 41 | + torch.set_default_tensor_type(torch.BFloat16Tensor) |
| 42 | + else: |
| 43 | + torch.set_default_tensor_type(torch.cuda.BFloat16Tensor) |
| 44 | + |
| 45 | + ckpt_paths = np.array(sorted(ckpt_paths)) |
| 46 | + |
| 47 | + new_mp_size, new_mp_rank = get_model_parallel_world_size(), get_model_parallel_rank() |
| 48 | + old_mp_size = len(ckpt_paths) |
| 49 | + old_mp_ranks = map_mp_rank(old_mp_size, new_mp_size, new_mp_rank) |
| 50 | + |
| 51 | + print(f"Loading checkpoint shards:\n{str(ckpt_paths[old_mp_ranks])}") # type: ignore |
| 52 | + paths = ckpt_paths[old_mp_ranks] # type: ignore |
| 53 | + state_dicts = [torch.load(str(p), map_location=map_location, mmap=mmap) for p in paths] |
| 54 | + |
| 55 | + if new_mp_size == old_mp_size: |
| 56 | + return state_dicts[0] # type: ignore |
| 57 | + |
| 58 | + if moe_num_experts is not None: |
| 59 | + state_dicts = [convert_moe_weights(d, moe_num_experts) for d in state_dicts] |
| 60 | + |
| 61 | + print(f"Resharding {len(state_dicts)} state dicts from MP size {old_mp_size} to MP size {new_mp_size}") |
| 62 | + return reshard_mp( |
| 63 | + state_dicts, |
| 64 | + size=max(new_mp_size // old_mp_size, 1), |
| 65 | + rank=new_mp_rank % max(new_mp_size // old_mp_size, 1), |
| 66 | + repeat_qk_qv=max(new_mp_size // n_kv_heads, 1), |
| 67 | + ) |
| 68 | + |
| 69 | + |
| 70 | +_WEIGHT_ROW_KEY = { |
| 71 | + "feed_forward.w2", |
| 72 | + "feed_forward.mlp.fc2", |
| 73 | + "attention.wo", |
| 74 | + "feed_forward.mlp.fc2_weight", |
| 75 | + "feed_forward.w_out_shared_DF.weight", |
| 76 | + "attn.wo.weight", |
| 77 | + "mlp.c_proj.weight", |
| 78 | +} |
| 79 | +_MOE_WEIGHT_ROW_KEY = {"feed_forward.experts.(moe_w_in_eD_F|moe_w_swiglu_eD_F)"} |
| 80 | + |
| 81 | +_WEIGHT_COLUMN_KEY = { |
| 82 | + "output", |
| 83 | + "feed_forward.(w1|w3)", |
| 84 | + "feed_forward.mlp.(fc1|fc3)", |
| 85 | + "feed_forward.mlp.fc1_weight", |
| 86 | + "attention.(wk|wq|wv|wqkv).weight", |
| 87 | + "feed_forward.(w_in_shared_FD|w_swiglu_FD)", |
| 88 | + "attn.(wk|wq|wv).weight", |
| 89 | + "attn.(wk|wq|wv).bias", |
| 90 | + "mlp.c_fc.weight", |
| 91 | + "mlp.c_fc.bias", |
| 92 | + "conv1._linear.weight", |
| 93 | + "tok_embeddings.weight", |
| 94 | + "vision_projection.weight", |
| 95 | +} |
| 96 | +_MOE_WEIGHT_COLUMN_KEY = {"feed_forward.experts.moe_w_out_eF_D"} |
| 97 | + |
| 98 | + |
| 99 | +def reshard_mp( |
| 100 | + state_dicts: List[Dict[str, torch.Tensor]], |
| 101 | + size: int, |
| 102 | + rank: int, |
| 103 | + repeat_qk_qv: int = 1, |
| 104 | +) -> Dict[str, torch.Tensor]: |
| 105 | + """ |
| 106 | + Reshard a list of state dicts into a single state dict given a change in MP size. |
| 107 | + If the list has more than one state dict, we concatenate the values of the same |
| 108 | + key across all state dicts. Otherwise, we just slice it for the current MP rank. |
| 109 | + """ |
| 110 | + |
| 111 | + def concat_or_chunk(tensors: List[torch.Tensor], dim: int) -> torch.Tensor: |
| 112 | + if len(tensors) > 1: |
| 113 | + return torch.cat(tensors, dim=dim) |
| 114 | + return tensors[0].chunk(size, dim=dim)[rank].clone() |
| 115 | + |
| 116 | + def process_key(key: str) -> torch.Tensor: |
| 117 | + if row_regex.search(key): |
| 118 | + return concat_or_chunk([s[key] for s in state_dicts], dim=-1) |
| 119 | + elif column_regex.search(key): |
| 120 | + if "w13" in key or "fc1_weight" in key: |
| 121 | + dims = state_dicts[0][key].size() |
| 122 | + values = [s[key].view(2, dims[0] // 2, *dims[1:]) for s in state_dicts] |
| 123 | + return concat_or_chunk(values, dim=1).flatten(0, 1) |
| 124 | + elif "qkv" in key: |
| 125 | + q_dim = state_dicts[0][key.replace("qkv", "o")].size(1) |
| 126 | + kv_dim = (state_dicts[0][key].size(0) - q_dim) // 2 |
| 127 | + values = [s[key].split((q_dim, kv_dim, kv_dim)) for s in state_dicts] |
| 128 | + return torch.cat([concat_or_chunk(x, dim=0) for x in zip(*values, strict=False)]) # type: ignore |
| 129 | + elif "wk.weight" in key or "wv.weight" in key: |
| 130 | + # Support MP > #kv_head |
| 131 | + return concat_or_chunk([s[key].repeat(repeat_qk_qv, 1) for s in state_dicts], dim=0) |
| 132 | + elif key == "output.bias" or key == "fc.weight": |
| 133 | + return concat_or_chunk([s[key] for s in state_dicts], dim=0) |
| 134 | + elif "w_" in key: |
| 135 | + return concat_or_chunk([s[key] for s in state_dicts], dim=-2) |
| 136 | + else: |
| 137 | + return concat_or_chunk([s[key] for s in state_dicts], dim=0) |
| 138 | + else: |
| 139 | + return state_dicts[0][key].clone() |
| 140 | + |
| 141 | + row_keys = _WEIGHT_ROW_KEY | _MOE_WEIGHT_ROW_KEY |
| 142 | + column_keys = _WEIGHT_COLUMN_KEY | _MOE_WEIGHT_COLUMN_KEY |
| 143 | + |
| 144 | + column_regex = re.compile("|".join(column_keys)) |
| 145 | + row_regex = re.compile("|".join(row_keys)) |
| 146 | + |
| 147 | + output: Dict[str, torch.Tensor] = {} |
| 148 | + with concurrent.futures.ThreadPoolExecutor() as executor: |
| 149 | + # Note: only processes keys in the first state dict. |
| 150 | + # Assumes keys are the same across all state dicts. |
| 151 | + mappings = {executor.submit(process_key, key): key for key in state_dicts[0]} |
| 152 | + for future in concurrent.futures.as_completed(mappings): |
| 153 | + output[mappings[future]] = future.result() |
| 154 | + return output |
| 155 | + |
| 156 | + |
| 157 | +def convert_moe_weights(state_dict: Dict[str, Any], num_experts: int) -> Dict[str, Any]: |
| 158 | + routed_keys = _MOE_WEIGHT_ROW_KEY | _MOE_WEIGHT_COLUMN_KEY |
| 159 | + routed_regex = re.compile("|".join(routed_keys)) |
| 160 | + keys = list(state_dict.keys()) |
| 161 | + for key in keys: |
| 162 | + if routed_regex.search(key): |
| 163 | + state_dict[key] = state_dict.pop(key).unflatten(0, (num_experts, -1)).squeeze(dim=0) |
| 164 | + return state_dict |
0 commit comments