待翻译：The whole of PyTorch on one page

AI 服务暂时不可用，以下为来源正文，待恢复后补全翻译。

Table of Contents The fall The territory The twelve ideas How this series draws How to read this What you can now say Try it yourself Pick a door References Figure 1. the program this whole series is about. You have typed something like this a thousand times. This series exists so that, by its end, you know everything these lines do. All of it: the Python they touch, the C++ they land in, the graph they record, the kernels they choose, the memory they use, and the two clocks they run on. Each of those words gets a plain meaning on its floor below. This is Part 0, the map. First we go down through all the layers once, fast. Then we draw the territory. Then twelve ideas that make the rest of the codebase predictable. Then how this series works, and how to read it. Nothing here gets its full story. Everything here gets a place, and every full story has a numbered part waiting for it. One promise before we start. Every measured number in this series comes from a small script you can run yourself, linked right where the number appears. I measured these on an Apple M3 Max laptop with torch 2.11.0 [1]. Your numbers will differ. The pattern they make will not. The fall PyTorch is deep. Between your keyboard and the chip there are eight levels. I will call them floors, and this meter shows all of them. It returns through the whole series, so you always know how deep you are. Figure 2. the depth meter. the orange dot marks where you are. The fastest way to learn a building is to go down through it once without stopping. That is this section. Floor one: python torch.randn looks like a Python function. Ask Python what it actually is:>>> type(torch.randn) Python gives that type only to functions written in compiled code. Compiled code means: code that was translated to machine instructions before you ever installed it, so there is no Python body inside it to read, and no line for your debugger to stop on. So where do those machine instructions live? In shared libraries. A shared library is a file of compiled code that a program loads while it runs. They sit inside the torch package on your disk, and you can look at them (proof):torch._C -> _C.cpython-312-darwin.so (49 KB, the loader) libtorch_cpu.dylib 206.5 MB (tensors and kernels) libtorch_python.dylib 28.5 MB (the python side of the border) p0_the_library.py the proof, ready to read or run """Proof: where the compiled part of pytorch actually lives. torch._C is a thin compiled stub; the weight of the framework is in the shared libraries next to it. Prints the files and their sizes. """ import glob import os import torch stub = torch._C.file print(f"torch {torch.version}") print(f"torch._C -> {os.path.basename(stub)} " f"({os.path.getsize(stub)/1024:.0f} KB stub)") libdir = os.path.join(os.path.dirname(stub), "lib") for lib in ["libtorch_cpu.dylib", "libtorch_python.dylib"]: p = os.path.join(libdir, lib) if os.path.exists(p): print(f"{lib:24s} {os.path.getsize(p)/1024/1024:6.1f} MB") download and run it Read the sizes, and then look at them: Figure 3. drawn to scale by file size. the part of pytorch that python can see is the orange dot. The part of PyTorch you can see from Python is a 49 KB file whose only job is to load the other two. The real body is 235 MB of compiled code. import torch brings it into your process, and after that, calling torch.randn means jumping into that body. Today we only need to know these files exist. This is the first honest surprise of the codebase: the Python you write all day is the smallest layer of it. The boundary The call leaves Python at once. Where does it land? In a C++ function named THPVariable_randn, inside that 28.5 MB library from the last floor. And here is a strange fact you can keep: this function does not exist in the PyTorch repository. Clone the repo, search for the name, and you find nothing. A program writes this function during the build, together with thousands of its siblings. Idea 4 below explains why, and Part 5 shows the program that does the writing. Figure 4. the border between the two languages. every tensor operation crosses it. Crossing this border costs time. To see the cost alone, time the smallest possible operation, where almost no arithmetic hides it (proof):add, 1 element : 0.538 microseconds per call add, 4M elements : 337.264 microseconds per call p2_dispatch_cost.py the proof, ready to read or run """Proof: the fixed cost of one eager op, and why size hides it. Times the same a + b at two sizes. The one-element add is nearly pure machinery (dispatch, wrapping, allocation); the 4M-element add is nearly pure arithmetic. CPU, single process. """ import time import torch def per_op_us(a, b, iters): # warmup for _ in range(2000): a + b t0 = time.perf_counter() for _ in range(iters): a + b return (time.perf_counter() - t0) / iters * 1e6 tiny = per_op_us(torch.ones(1), torch.ones(1), 200_000) big_n = 4_000_000 big = per_op_us(torch.ones(big_n), torch.ones(big_n), 2_000) print(f"torch {torch.version}, cpu") print(f"add, 1 element : {tiny:8.3f} us/op") print(f"add, 4M elements : {big:8.3f} us/op") print(f"machinery share of the tiny op: ~all of it") print(f"ops/sec you can issue from python: {1e6/tiny:,.0f}") # The sweep behind the toll meter: the same add at twelve sizes, # 1 to 4M elements in powers of four. Every dot on the widget's # axis is one line of this output. import json sweep = [] for k in range(12): n = 4 ** k iters = max(1_000, min(200_000, 40_000_000 // max(n, 1))) us = per_op_us(torch.ones(n), torch.ones(n), iters) sweep.append({"n": n, "us": round(us, 3)}) print(f"add, {n:>9,} elements : {us:9.3f} us/op") print("JSON_SWEEP=" + json.dumps(sweep)) download and run it The one-element add does almost no math. So its 0.54 microseconds is almost pure crossing cost: leave Python, check the arguments, build the result object, return. Half a microsecond sounds like nothing. It means Python can issue at most about 1.9 million operations per second, and a single training step contains thousands of operations. Keep this number. It returns in Idea 6. The dispatcher Under the border, the call reaches the strangest machine in PyTorch: the dispatcher. The dispatcher is the router that decides, for every operation, which pieces of code run and in what order. Look at what it must decide. Your three lines never said "record gradients". No if statement in your code turns that on. Yet somewhere, something decided that this matrix multiplication should be remembered for backward(). That something is the dispatcher. Every operation passes down through a fixed stack of layers. Each layer can act on the call, change it, or let it pass unchanged. Autograd, the part of PyTorch that computes gradients, is one such layer. Mixed precision is another. On this run, only autograd is awake. Figure 5. four layers touch your call before any arithmetic starts. only the highlighted one is awake today. The kernel At the bottom of the stack, one concrete function is chosen. Chosen is the right word. This torch build has 3,677 registered operation names (proof prints the count), and a name is not a function body. The operation addmm, the matrix multiplication behind model(x), has separate bodies for CPU and for each kind of GPU, for each data type, for dense and for sparse tensors. A body like this, written for one device and one data type, is called a kernel. The dispatcher's last job is to pick one: Figure 6. one name, a grid of bodies. the dispatcher picks exactly one cell per call. p4_micro_proofs.py the proof, ready to read or run """Micro-proofs quoted in Part 0: storage sharing, view errors, mutation rewriting history, the no_grad layer, float32 absorption.""" import torch print(f"torch {torch.version}\n") # 1. a tensor is a window over storage x = torch.arange(6, dtype=torch.float32) v = x.view(2, 3) print("same bytes under both:", x.data_ptr() == v.data_ptr()) print("v.stride():", v.stride(), " v.t().stride():", v.t().stride()) try: v.t().view(-1) except RuntimeError as e: print("v.t().view(-1) ->", str(e).split(".")[0]) # 2. mutation rewrites the recorded program a = torch.ones(3, requires_grad=True) y = a * 2 print("\nbefore add_:", type(y.grad_fn).name) y.add_(1) print("after add_:", type(y.grad_fn).name) # 3. no_grad removes one dispatcher layer with torch.no_grad(): z = a * 2 print("\ninside no_grad, grad_fn:", z.grad_fn) # 4. float32 absorbs small numbers t = torch.tensor(1e8) print("\n(1e8 + 1) - 1e8 in float32 =", ((t + 1) - t).item()) # 5. the size of the operation list (idea 3) print("\nregistered operation names:", len(torch._C._dispatch_get_all_op_names())) download and run it The full list of operations lives in one file in the repository: native_functions.yaml [2]. Its sibling derivatives.yaml [3] lists the derivative of each operation. Everything else grows from these two files. No other file in the repository tells you as much per line. Figure 7. 3,677 names on the left. one function body on the right. the funnel is the dispatcher's last job. One floor down sits memory. torch.randn(64, 128) needs 32,768 bytes: 64 rows times 128 numbers times 4 bytes per number. On the CPU this is an ordinary allocation. On a GPU it is not. There, PyTorch runs its own allocator, a keeper of memory that asks the GPU driver for large blocks once and then reuses them, because asking the driver every time is slow. This allocator decides when you run out of memory and what the error means. Part 4 examines it. The two clocks Here the story splits in two. What follows is the single most useful performance fact in PyTorch. On a GPU, your Python line does not do the work. It requests the work, and the request returns at once. The GPU does the work on its own clock, while Python continues. I measured it on this machine's GPU (proof):time to request 50 matrix multiplications : 1.58 ms time until the work was actually done : 73.83 ms python was free during : 72.25 ms (98%) p3_two_timelines.py the proof, ready to read or run """Proof: the CPU runs ahead of the GPU. Queues 50 large matmuls on the MPS device and measures two times: how long Python took to *ask* for the work, and how long the work actually took. The difference is the gap the chapter draws. """ import time import torch assert torch.backends.mps.is_available(), "needs an Apple-silicon GPU" dev = torch.device("mps") a = torch.randn(2048, 2048, device=dev) b = torch.randn(2048, 2048, device=dev) for _ in range(5): # warmup (a @ b) torch.mps.synchronize() t0 = time.perf_counter() for _ in range(50): c = a @ b t_queue = time.perf_counter() - t0 torch.mps.synchronize() t_done = time.perf_counter() - t0 print(f"torch {torch.version}, mps") print(f"time to queue 50 matmuls : {t_queue*1e3:8.2f} ms") print(f"time until work finished : {t_done*1e3:8.2f} ms") print(f"python was free for : {(t_done-t_queue)*1e3:8.2f} ms ({(t_done-t_queue)/t_done:.0%} of the wall time)") # The three ways to read the loss, measured, for the two-clocks # widget: never, once at the end, after every step. import json def run_mode(mode, iters=50): for _ in range(5): (a @ b) torch.mps.synchronize() t0 = time.perf_counter() t_free = 0.0 for i in range(iters): c = a @ b if mode == "every": c[0, 0].item() t_q = time.perf_counter() - t0 if mode == "once": c[0, 0].item() torch.mps.synchronize() total = time.perf_counter() - t0 return {"mode": mode, "queue_ms": round(t_q * 1e3, 2), "total_ms": round(total * 1e3, 2), "free_ms": round((total - t_q) * 1e3, 2)} modes = [run_mode(m) for m in ("never", "once", "every")] for m in modes: print(f"read {m['mode']:>5}: total {m['total_ms']:8.2f} ms, " f"python busy {m['queue_ms']:8.2f} ms") print("JSON_MODES=" + json.dumps(modes)) download and run it Figure 8. two clocks, one program. the cpu requested everything in the first two milliseconds; the gpu needed seventy-two more to finish. Python asked for all f [truncated for AI cost control]