CodeVetter · Repo Unpack · Data flow
pytorch/pytorch
The core data path is: Python call -> torch.* binding (codegen'd _VariableFunctions) or torch.ops namespace -> C++ dispatcher keyed by DispatchKey -> backend kernel (CPU/CUDA/etc.) -> TensorImpl/Storage in c10/ATen, with autograd Nodes attached during forward when grad is required. Compile mode inserts a bytecode interception layer: Dynamo's PEP 523 hook traces Python frames into FX graphs, AOT Autograd splits forward/backward, and Inductor lowers to Triton/C++ kernels. Persistence flows through a custom zip format where storages are extracted via pickle persistent_id and written as raw byte records. Distributed state flows through Store-based rendezvous into backend ProcessGroups.
- torch.ops.<ns>.<op> lazily resolves through _OpNamespace.__getattr__ -> OpOverloadPacket.__getattr__ -> torch._C._get_operation_overload + _get_schema -> cached OpOverload; calling it invokes self._op (the C++ unboxed dispatcher entry).
- OpOverload keeps a per-DispatchKey _dispatch_cache and a py_kernels table; decompose() prefers a Python CompositeImplicitAutograd kernel and otherwise redispatches into C++ via _op_dk.
- native_functions.yaml is the schema source of truth: each entry is `func: name(ArgType...) -> Return` plus a `dispatch:` block mapping backend keys (CPU, CUDA, ...) to kernel function names; torchgen parse_native_yaml turns these into NativeFunctions + backend indices that drive all generated bindings.
- Autograd data flow: forward ops record grad_fn Nodes; backward() -> _engine_run_backward stashes the contextvars context in TLS then calls Variable._execution_engine.run_backward; C++ Engine::execute validates outputs, creates a GraphTask and schedules NodeTasks on per-device worker threads.
- torch.save path: save() -> _save() runs a PyTorchPickler whose persistent_id intercepts storages, dedupes them by _cdata into serialized_storages, enforces dtype-consistency per data_ptr, writes data.pkl + .format_version + .storage_alignment + byteorder records, then streams each storage (moved to CPU, optionally via pinned-memory D2H) as data/<key>.
- torch.load path: stricter _is_zipfile check -> _open_zipfile_reader (torch._C.PyTorchFileReader) -> _load() reads byteorder/.format_version/.storage_alignment, computes restore_location from map_location, and load_tensor() builds UntypedStorages (meta storages when a fake mode is active, recording _checkpoint_offset for layout-preserving reload).
- weights_only=True is the default when pickle_module is unset; _weights_only_unpickler restricts pickle GLOBALs to an allowlist extensible via add_safe_globals()/safe_globals(), and get_unsafe_globals_in_checkpoint scans data.pkl for non-allowlisted globals.
- torch.compile data flow: compile() -> torch._dynamo.optimize -> _optimize -> OptimizeContext wrapping the fn; the C++ PEP 523 eval-frame hook intercepts frames, symbolic_convert builds an FX graph in OutputGraph with guards, then the backend (compile_fx) is invoked and generated bytecode calls the compiled artifact with resume functions after graph breaks.
- Inductor's compile_fx orchestrates end-to-end compilation: selects the decomposition table, wakes the AsyncCompile worker pool, wraps inner_compile into fw_compiler/bw_compiler callables, and calls aot_autograd which partitions the graph and invokes those compilers before fx_codegen_and_compile emits runnable code.
- init_process_group resolves the default backend (Backend(...) validation, device-based default), performs rendezvous through init_method/env:// to get a Store, then _new_process_group_helper constructs the ProcessGroup; all collective payloads traverse the backend with no auth or encryption.
- Compile caches are per code object, not per frame: multiple compiled entries accumulate per code object until torch._dynamo.config.recompile_limit (default 8), after which execution falls back to eager; isolate_recompiles gives a compile call its own entry set.
- Distributed checkpoint (torch.distributed.checkpoint) deliberately bypasses the weights_only protections: checkpoints are assumed to come from the user's own trusted store, so loading may execute arbitrary code.