QLoRA › Memory Savings Analysis

def calculate_memory_breakdown(
    n_params_billions,
    lora_rank=16,
    lora_target_frac=0.02,  # Fraction of params targeted by LoRA
    seq_len=512,
    batch_size=1,
    hidden_dim=None,
):
    """Calculate memory usage for different fine-tuning approaches."""
    n_params = n_params_billions * 1e9

    # Estimate hidden dimension from parameter count (rough approximation)
    if hidden_dim is None:
        hidden_dim = int(np.sqrt(n_params / 100))  # Very rough estimate

    n_lora_params = n_params * lora_target_frac * 2 * lora_rank / hidden_dim

    results = {}

    # Full fine-tuning in FP16
    results["Full FP16"] = {
        "model": n_params * 2,  # FP16 weights
        "gradients": n_params * 2,  # FP16 gradients
        "optimizer": n_params * 8,  # FP32 Adam states
        "total": n_params * 12,
    }

    # LoRA in FP16
    results["LoRA FP16"] = {
        "model": n_params * 2,  # FP16 base weights
        "gradients": n_lora_params * 2,  # FP16 LoRA gradients
        "optimizer": n_lora_params * 8,  # FP32 Adam states for LoRA
        "total": n_params * 2 + n_lora_params * 10,
    }

    # QLoRA
    nf4_bits_per_param = 4 + 0.125  # 4-bit weights + double-quantized scales
    results["QLoRA"] = {
        "model": n_params * nf4_bits_per_param / 8,  # NF4 base weights
        "gradients": n_lora_params * 2,  # BF16 LoRA gradients
        "optimizer": n_lora_params * 8,  # FP32 Adam states for LoRA
        "total": n_params * nf4_bits_per_param / 8 + n_lora_params * 10,
    }

    # Convert to GB
    for method in results:
        for key in results[method]:
            results[method][key] /= 1e9

    return results


# Calculate for different model sizes
model_sizes = [7, 13, 34, 70]
memory_data = {size: calculate_memory_breakdown(size) for size in model_sizes}