#!/usr/bin/env python3 """Minimal LoRA/QLoRA budget checks plus an optional PyTorch toy run.""" from __future__ import annotations import argparse import math import random from dataclasses import dataclass @dataclass(frozen=True) class LinearBudget: in_features: int out_features: int rank: int alpha: int @property def full_params(self) -> int: return self.in_features * self.out_features @property def lora_params(self) -> int: return self.rank * (self.in_features + self.out_features) @property def trainable_fraction(self) -> float: return self.lora_params / self.full_params @property def scaling(self) -> float: return self.alpha / self.rank def estimate_model_memory( params: int, trainable_params: int, *, base_bits: int = 16, trainable_bits: int = 16, optimizer_bits_per_param: int = 64, ) -> dict[str, float]: """Return rough memory components in MiB. This is intentionally a budget sketch, not a hardware profiler. It excludes activations, temporary kernels, allocator fragmentation, gradients for frozen weights, and framework metadata. """ frozen_params = max(params - trainable_params, 0) frozen_weight_mib = frozen_params * base_bits / 8 / 1024**2 trainable_weight_mib = trainable_params * trainable_bits / 8 / 1024**2 optimizer_mib = trainable_params * optimizer_bits_per_param / 8 / 1024**2 return { "frozen_weight_mib": frozen_weight_mib, "trainable_weight_mib": trainable_weight_mib, "optimizer_state_mib": optimizer_mib, "total_tracked_mib": frozen_weight_mib + trainable_weight_mib + optimizer_mib, } def print_budget(args: argparse.Namespace) -> None: budget = LinearBudget(args.in_features, args.out_features, args.rank, args.alpha) print("=== LoRA single linear layer budget ===") print(f"base_weight_shape=[{budget.out_features}, {budget.in_features}]") print(f"lora_A_shape=[{budget.rank}, {budget.in_features}]") print(f"lora_B_shape=[{budget.out_features}, {budget.rank}]") print(f"delta_W_shape=[{budget.out_features}, {budget.in_features}]") print(f"full_params={budget.full_params}") print(f"lora_trainable_params={budget.lora_params}") print(f"trainable_fraction={budget.trainable_fraction:.6f}") print(f"scaling_alpha_over_r={budget.scaling:.3f}") model_params = args.model_params full = estimate_model_memory(model_params, model_params, base_bits=16) lora = estimate_model_memory(model_params, args.lora_trainable_params, base_bits=16) qlora = estimate_model_memory(model_params, args.lora_trainable_params, base_bits=4) print("\n=== Rough model memory budget, MiB ===") print("scenario,frozen_weight,trainable_weight,optimizer_state,total_tracked") for name, item in [("full_fp16", full), ("lora_fp16_base", lora), ("qlora_4bit_base", qlora)]: print( f"{name},{item['frozen_weight_mib']:.1f},{item['trainable_weight_mib']:.1f}," f"{item['optimizer_state_mib']:.1f},{item['total_tracked_mib']:.1f}" ) def run_torch_toy(args: argparse.Namespace) -> None: try: import torch from torch import nn from torch.nn import functional as F except ModuleNotFoundError as exc: raise SystemExit("PyTorch is not installed; run --check-only for stdlib budget checks.") from exc torch.manual_seed(args.seed) random.seed(args.seed) class LoRALinear(nn.Module): def __init__(self, in_features: int, out_features: int, rank: int, alpha: int) -> None: super().__init__() self.base = nn.Linear(in_features, out_features, bias=True) for param in self.base.parameters(): param.requires_grad_(False) self.lora_A = nn.Parameter(torch.empty(rank, in_features)) self.lora_B = nn.Parameter(torch.zeros(out_features, rank)) nn.init.kaiming_uniform_(self.lora_A, a=math.sqrt(5)) self.scaling = alpha / rank def forward(self, x: torch.Tensor) -> torch.Tensor: delta = (x @ self.lora_A.t() @ self.lora_B.t()) * self.scaling return self.base(x) + delta model = LoRALinear(args.in_features, 2, args.rank, args.alpha) opt = torch.optim.AdamW([model.lora_A, model.lora_B], lr=args.lr) trainable = sum(p.numel() for p in model.parameters() if p.requires_grad) frozen = sum(p.numel() for p in model.parameters() if not p.requires_grad) print("=== PyTorch toy LoRA classifier ===") print(f"input_shape=[{args.batch_size}, {args.in_features}]") print(f"logit_shape=[{args.batch_size}, 2]") print(f"trainable_params={trainable}, frozen_params={frozen}") for step in range(1, args.steps + 1): x = torch.randn(args.batch_size, args.in_features) y = (x[:, 0] + 0.7 * x[:, 3] - 0.5 * x[:, 7] > 0).long() logits = model(x) loss = F.cross_entropy(logits, y) opt.zero_grad() loss.backward() opt.step() if step in {1, args.steps}: pred = logits.argmax(dim=-1) acc = (pred == y).float().mean().item() print(f"step={step}, loss={loss.item():.4f}, batch_acc={acc:.3f}") def main() -> None: parser = argparse.ArgumentParser() parser.add_argument("--in-features", type=int, default=64) parser.add_argument("--out-features", type=int, default=64) parser.add_argument("--rank", type=int, default=4) parser.add_argument("--alpha", type=int, default=16) parser.add_argument("--model-params", type=int, default=7_000_000_000) parser.add_argument("--lora-trainable-params", type=int, default=20_000_000) parser.add_argument("--check-only", action="store_true") parser.add_argument("--steps", type=int, default=40) parser.add_argument("--batch-size", type=int, default=64) parser.add_argument("--lr", type=float, default=3e-2) parser.add_argument("--seed", type=int, default=7) args = parser.parse_args() if args.rank <= 0: raise SystemExit("--rank must be positive") if args.alpha <= 0: raise SystemExit("--alpha must be positive") print_budget(args) if not args.check_only: run_torch_toy(args) if __name__ == "__main__": main()