Calcular VRAM necessária para Model

calc.py

def calcular_vram(num_parametros, quantization_bits, overhead_factor=1.2):
    memoria_modelo = (num_parametros * quantization_bits) / (8 * 10**9)
    memoria_total = memoria_modelo * overhead_factor
    return memoria_total


def main():
    try:
        num_parametros = float(input("Número de parâmetros do modelo (em bilhões, ex: 8 para 8B): ")) * 10**9
        quantization_bits = int(input("Tamanho da quantização (quantization size em bits, ex: 4, 6, 8): "))
        overhead_input = input("Informe o fator de overhead (Padrão 1.2): ")

        if overhead_input == "":
            overhead_factor = 1.2
        else:
            overhead_factor = float(overhead_input)
            if overhead_factor == 0:
                overhead_factor = 1

        vram_utilizada = calcular_vram(num_parametros, quantization_bits, overhead_factor)

        print(f"\nVRAM necessária é: {vram_utilizada:.2f} GB")

    except ValueError:
        print("Insira números válidos.")


if __name__ == "__main__":
    main()