YangWang92

import torch
from argparse import ArgumentParser
from vptq.utils.hessian import load_hessian
import os

if __name__ == "__main__":
    parser = ArgumentParser()
    parser.add_argument('--load_hessian_dir', type=str, default=None,
                        help='Directory containing Hessian .pt files')
    parser.add_argument('--store_inv_hessian_dir', type=str, default=None,

YangWang92

import torch

def merge_model_parts(part_files, output_file, merge_dict):
    merged_state_dict = {}
    _merged_state_dict = {}
    for part_file in part_files:
        print(f'load part {part_file}')
        state_dict = torch.load(part_file)
        for name, param in state_dict.items():
            print(f'load {name}')

YangWang92

# script to decompose/recompose llama model in different number of shards
# note that it loads the full model * 2 in cpu memory

import os
import json
import sys
import torch
import glob

if len(sys.argv) != 4:

YangWang92

javascript:void(document.documentElement.style.filter='invert(65%) contrast(370%) brightness(220%) hue-rotate(110deg)')

YangWang92

import inspect
import logging
import os
from pathlib import Path

import torch
from psutil import cpu_count
from transformers import T5Config, T5ForConditionalGeneration, T5Tokenizer
from transformers.generation_utils import GenerationMixin
from transformers.modeling_outputs import BaseModelOutputWithPast, Seq2SeqLMOutput