def save_shards(model_sd, num_shards: int, prefix="", verbose=False):
    """
    Convert and save the HF format weights to PTH format weights
    """
    with torch.no_grad():
        if num_shards == 1:
            new_state_dict = {}
            for k, v in model_sd.items():
                new_k = translate_state_dict_key(k)
                if new_k is not None:
                    if "wq" in new_k or "wk" in new_k:
                        new_state_dict[new_k] = unpermute(v)
                    else:
                        new_state_dict[new_k] = v

            os.makedirs(output_dir, exist_ok=True)
            print(
                f"Saving shard 1 of {num_shards} into"
                f" {output_dir}/{prefix}consolidated.00.pth"
            )
            torch.save(new_state_dict, output_dir + f"/{prefix}consolidated.00.pth")
        else:
            new_state_dicts = [dict() for _ in range(num_shards)]
            for k in list(model_sd.keys()):
                v = model_sd[k]
                new_k = translate_state_dict_key(k)
                if new_k is not None:
                    if new_k == "tok_embeddings.weight":
                        assert v.size(1) % num_shards == 0
                        splits = v.split(v.size(1) // num_shards, dim=1)
                    elif new_k == "output.weight":
                        if v.size(0) % num_shards == 0:
                            splits = v.split(v.size(0) // num_shards, dim=0)
                        else:
                            size_list = [v.size(0) // num_shards] * num_shards
                            size_list[-1] += v.size(0) % num_shards
                            splits = v.split(
                                size_list, dim=0
                            )  # 13B: size_list == [24976,24977]
                    elif new_k == "norm.weight":
                        splits = [v] * num_shards
                    elif "ffn_norm.weight" in new_k:
                        splits = [v] * num_shards
                    elif "attention_norm.weight" in new_k:
                        splits = [v] * num_shards

                    elif "w1.weight" in new_k:
                        splits = v.split(v.size(0) // num_shards, dim=0)
                    elif "w2.weight" in new_k:
                        splits = v.split(v.size(1) // num_shards, dim=1)
                    elif "w3.weight" in new_k:
                        splits = v.split(v.size(0) // num_shards, dim=0)

                    elif "wo.weight" in new_k:
                        splits = v.split(v.size(1) // num_shards, dim=1)

                    elif "wv.weight" in new_k:
                        splits = v.split(v.size(0) // num_shards, dim=0)

                    elif "wq.weight" in new_k or "wk.weight" in new_k:
                        v = unpermute(v)
                        splits = v.split(v.size(0) // num_shards, dim=0)
                    else:
                        print(f"Unexpected key {new_k}")
                        raise ValueError
                    if verbose:
                        print(f"Processing {new_k}")
                    for sd, split in zip(new_state_dicts, splits):
                        sd[new_k] = split.clone()
                        del split
                    del splits
                del model_sd[k], v
                gc.collect()  # Effectively enforce garbage collection

            os.makedirs(output_dir, exist_ok=True)
            for i, new_state_dict in enumerate(new_state_dicts):
                print(
                    f"Saving shard {i+1} of {num_shards} into"
                    f" {output_dir}/{prefix}consolidated.0{i}.pth"
                )
                torch.save(
                    new_state_dict, output_dir + f"/{prefix}consolidated.0{i}.pth"
                )

    Convert and save the HF format weights to PTH format weights
    
smoe/utils/merge_llama_with_lora.py