어떤 일에 어떤 기계

#!/usr/bin/env python3
"""Which machine for which job -- the decision the physics track licenses, as a function.
Inputs: the bytes a token READS (GB, at your bits), the bytes that must be RESIDENT (GB;
the same number for a dense model, the whole file for a mixture), how many streams at
once, the parameter count if you train, and whether the box must sit in a room with people.
Every rule below is arithmetic from earlier lessons; the ranking is this quest's judgment."""

MACHINES = {  # name: (memory GB, GB/s, watts max, quiet in a room?)
    "MacBook Pro M5 Max 128 GB":      (128,  614,  140, True),
    "Mac Studio M3 Ultra 512 GB":     (512,  819,  270, True),
    "RTX 5090 workstation":           (32,   1792, 1000, False),
    "RTX PRO 6000 workstation":       (96,   1792, 1000, False),
    "DGX Spark":                      (128,  273,  140, True),
    "rented H100 (cloud)":            (80,   3350, 700, True),      # quiet because it is elsewhere
}


def decide(read_gb, resident_gb=None, streams=1, train_params_b=None, in_room=True):
    resident_gb = read_gb if resident_gb is None else resident_gb
    rows = []
    for name, (mem, bw, watts, quiet) in MACHINES.items():
        fits = resident_gb + 4 <= mem * 0.9
        ceiling = bw / read_gb if fits else 63 / read_gb          # spill: PCIe 5 x16 carries the spilled bytes
        note = []
        if not fits: note.append("spills over PCIe" if "RTX" in name else "does not fit")
        if train_params_b and mem < train_params_b * 16:          # 16 bytes per parameter before activations
            note.append(f"no full fine-tune ({train_params_b}B x 16 B/param = {train_params_b * 16:.0f} GB)")
        if in_room and not quiet: note.append("loud/hot in a room")
        if streams > 1: note.append(f"batch x{streams}: compute-bound, favours tensor cores")
        rows.append((name, fits, ceiling, note))
    rows.sort(key=lambda r: (r[1], r[2]), reverse=True)
    print(f"reads {read_gb:.1f} GB/token, resident {resident_gb:.0f} GB, streams {streams}, "
          f"train={train_params_b or '-'}, in a room={in_room}")
    for name, fits, ceiling, note in rows:
        print(f"  {name:30} {'fits ' if fits else 'NO   '} ceiling {ceiling:6.1f} tok/s   {'; '.join(note)}")
    print()


decide(14.4)                                   # a 27B at 4 bits, one user
decide(39.7)                                   # a 70B at 4.5 bits/weight, one user
decide(12.4, resident_gb=132)                  # a 235B mixture, 22B active, at 4.5 bits: all experts resident, 22B read
decide(14.4, streams=32)                       # serving 32 users a 27B
decide(4.5, train_params_b=9)                  # fine-tuning a 9B for real

# reads 39.7 GB/token (a 70B at 4.5 bits/weight), one stream, in a room:
#   rented H100 (cloud)            fits  ceiling   84.4 tok/s
#   RTX PRO 6000 workstation       fits  ceiling   45.1 tok/s   loud/hot in a room
#   Mac Studio M3 Ultra 512 GB     fits  ceiling   20.6 tok/s
#   MacBook Pro M5 Max 128 GB      fits  ceiling   15.5 tok/s
#   DGX Spark                      fits  ceiling    6.9 tok/s
#   RTX 5090 workstation           NO    ceiling    1.6 tok/s   spills over PCIe; loud/hot in a room
# reads 12.4 GB/token, resident 132 GB (a 235B mixture, 22B active): only the 512 GB Studio fits, ceiling 66 tok/s
#   by the bytes a token reads -- dividing by the whole 132 GB would say 6, and that is the dense mistake
# reads 4.5 GB/token, train 9B: every card, the Spark and the 128 GB laptop say 'no full fine-tune (144 GB)'; the Studio does not