Skip to content

Commit

Permalink
Switch to using total VRAM instead of free VRAM to estimate tile size (
Browse files Browse the repository at this point in the history
…#2929)

* Switch to using total VRAM instead of free VRAM to estimate tile size

* use better values
  • Loading branch information
joeyballentine authored Jun 2, 2024
1 parent 7047b95 commit 8caf78b
Showing 1 changed file with 6 additions and 3 deletions.
Original file line number Diff line number Diff line change
Expand Up @@ -68,10 +68,13 @@ def estimate():
if options.use_fp16:
model_bytes = model_bytes // 2
mem_info: tuple[int, int] = torch.cuda.mem_get_info(device) # type: ignore
free, _total = mem_info
_free, total = mem_info
# only use 75% of the total memory
total = int(total * 0.75)
if options.budget_limit > 0:
free = min(options.budget_limit * 1024**3, free)
budget = int(free * 0.8)
total = min(options.budget_limit * 1024**3, total)
# Estimate using 80% of the value to be more conservative
budget = int(total * 0.8)

return MaxTileSize(
estimate_tile_size(
Expand Down

0 comments on commit 8caf78b

Please sign in to comment.