hf-accelerate/model-memory-usage

1# Utilities related to loading in and working with models/specific models2from urllib.parse import urlparse3 4import gradio as gr5import torch6from accelerate.commands.estimate import check_has_model, create_empty_model, estimate_training_usage7from accelerate.utils import calculate_maximum_sizes, convert_bytes8from huggingface_hub.utils import GatedRepoError, RepositoryNotFoundError9 10 11DTYPE_MODIFIER = {"float32": 1, "float16/bfloat16": 2, "int8": 4, "int4": 8}12 13 14def extract_from_url(name: str):15    "Checks if `name` is a URL, and if so converts it to a model name"16    is_url = False17    try:18        result = urlparse(name)19        is_url = all([result.scheme, result.netloc])20    except Exception:21        is_url = False22    # Pass through if not a URL23    if not is_url:24        return name25    else:26        path = result.path27        return path[1:]28 29 30def translate_llama(text):31    "Translates Llama-2 and CodeLlama to its hf counterpart"32    if not text.endswith("-hf"):33        return text + "-hf"34    return text35 36 37def get_model(model_name: str, library: str, access_token: str):38    "Finds and grabs model from the Hub, and initializes on `meta`"39    if "meta-llama/Llama-2-" in model_name or "meta-llama/CodeLlama-" in model_name:40        model_name = translate_llama(model_name)41    if library == "auto":42        library = None43    model_name = extract_from_url(model_name)44    try:45        model = create_empty_model(model_name, library_name=library, trust_remote_code=True, access_token=access_token)46    except GatedRepoError:47        raise gr.Error(48            f"Model `{model_name}` is a gated model, please ensure to pass in your access token and try again if you have access. You can find your access token here : https://huggingface.co/settings/tokens. "49        )50    except RepositoryNotFoundError:51        raise gr.Error(f"Model `{model_name}` was not found on the Hub, please try another model name.")52    except ValueError:53        raise gr.Error(54            f"Model `{model_name}` does not have any library metadata on the Hub, please manually select a library_name to use (such as `transformers`)"55        )56    except (RuntimeError, OSError) as e:57        library = check_has_model(e)58        if library != "unknown":59            raise gr.Error(60                f"Tried to load `{model_name}` with `{library}` but a possible model to load was not found inside the repo."61            )62        raise gr.Error(63            f"Model `{model_name}` had an error, please open a discussion on the model's page with the error message and name: `{e}`"64        )65    except ImportError:66        # hacky way to check if it works with `trust_remote_code=False`67        model = create_empty_model(68            model_name, library_name=library, trust_remote_code=False, access_token=access_token69        )70    except Exception as e:71        raise gr.Error(72            f"Model `{model_name}` had an error, please open a discussion on the model's page with the error message and name: `{e}`"73        )74    return model75 76 77def calculate_memory(model: torch.nn.Module, options: list):78    "Calculates the memory usage for a model init on `meta` device"79    total_size, largest_layer = calculate_maximum_sizes(model)80 81    data = []82    for dtype in options:83        dtype_total_size = total_size84        dtype_largest_layer = largest_layer[0]85 86        modifier = DTYPE_MODIFIER[dtype]87        dtype_training_size = estimate_training_usage(88            dtype_total_size, dtype if dtype != "float16/bfloat16" else "float16"89        )90        dtype_total_size /= modifier91        dtype_largest_layer /= modifier92 93        dtype_total_size = convert_bytes(dtype_total_size)94        dtype_largest_layer = convert_bytes(dtype_largest_layer)95        data.append(96            {97                "dtype": dtype,98                "Largest Layer or Residual Group": dtype_largest_layer,99                "Total Size": dtype_total_size,100                "Training using Adam (Peak vRAM)": dtype_training_size,101            }102        )103    return data104