Enhance documentation and update .gitignore for model conversion scripts

This commit is contained in:
enoch kan
2025-01-05 18:18:18 +00:00
parent fd011c11aa
commit a1296f099e
5 changed files with 179 additions and 1 deletions

View File

@@ -31,6 +31,18 @@ mapping = {
def main(hf_ckpt_path, save_path, n_experts, mp):
"""
Converts and saves model checkpoint files into a specified format.
Args:
hf_ckpt_path (str): Path to the directory containing the input checkpoint files.
save_path (str): Path to the directory where the converted checkpoint files will be saved.
n_experts (int): Total number of experts in the model.
mp (int): Model parallelism factor.
Returns:
None
"""
torch.set_num_threads(8)
n_local_experts = n_experts // mp
state_dicts = [{} for _ in range(mp)]