CodeVetter · Repo Unpack · Behavior traces
pytorch/pytorch
Traced flows: (1) import torch walks native-library preloading before any Python API exists; (2) torch.compile goes Python wrapper -> dynamo OptimizeContext -> C++ frame hook -> FX graph -> inductor -> aot_autograd -> Triton codegen -> guard-cached bytecode; (3) torch.save/load round-trips through a pickle-metadata + raw-storage zip format with the restricted unpickler on the read side; (4) backward() descends from torch.autograd into the C++ Engine GraphTask scheduler; (5) init_process_group does store rendezvous then backend construction; (6) the test runner shards test files and invokes pytest per module.
- import torch order: _rocm_init.initialize() if present -> (Windows) _load_dll_libraries / (POSIX) _load_global_deps ctypes.CDLL(libtorch_global_deps, RTLD_GLOBAL) + optional _preload_cuda_deps -> `from torch._C import *` -> _C._initExtension(torch_shm_manager path) -> copy _C._VariableFunctions into globals -> import subpackages (cpu/cuda/xpu/mps/mtia before multiprocessing) -> _C._init_names(storage classes) -> attach docstrings -> bind torch.ops/torch.classes namespaces.
- torch.compile call sequence: compile() logs API usage, rejects Python >=3.16 / free-threaded <3.13.3, resolves backend via get_default_backend, wraps inductor in _TorchCompileInductorWrapper, then torch._dynamo.optimize -> _optimize -> check_if_dynamo_supported -> (TORCHDYNAMO_DISABLE/justknobs check) -> optimize_assert if nopython else _optimize_catch_errors(convert_frame.convert_frame(backend,...)) -> OptimizeContext; per-call the C++ eval-frame hook runs guards then either executes cached compiled code or triggers compilation.
- Inside a dynamo compile: InstructionTranslator in symbolic_convert.py executes bytecode into VariableTrackers; OutputGraph owns the FX graph/guards; compile_subgraph calls the backend; PyCodegen emits replacement bytecode; resume_execution.py builds continuation functions for graph breaks; CheckFunctionManager builds a C++ GuardManager tree for fast recheck.
- torch.save sequence: save() -> _open_zipfile_writer -> _save() -> PyTorchPickler(persistent_id) dumps obj to data.pkl -> write .format_version/.storage_alignment/byteorder records -> for each serialized storage, move to CPU (pinned-memory D2H if configured) -> zip_file.write_record('data/<key>', storage, num_bytes).
- torch.load sequence: torch._C._log_api_usage_once('torch.load') -> default weights_only=True -> _is_zipfile -> _open_zipfile_reader -> _load(): read byteorder/format_version/alignment -> Unpickler (weights-only restricted unpickler or full pickle module) resolves 'storage' persistent ids via load_tensor -> restore_location applies map_location -> storage rebuilt from data/<key> record (or lazily via mmap).
- loss.backward() sequence: torch.autograd.backward -> functorch guard -> grad_variables deprecation shim -> __torch_function__ override check -> _tensor_or_tensors_to_tuple + _make_grads -> retain_graph defaults to create_graph -> _engine_run_backward (TLS context stash, optional logging hooks) -> Variable._execution_engine.run_backward -> C++ Engine::execute -> validate_outputs -> GraphTask scheduled on worker threads -> .grad accumulated on leaves (accumulate_grad=True).
- init_process_group sequence: reject double init (GroupMember.WORLD check) -> set_pytorch_distributed_envs_from_justknobs -> Backend(backend) validation/defaulting -> store resolution via explicit store or rendezvous(init_method, default env://) -> _new_process_group_helper builds the backend ProcessGroup (eager ncclCommInit* when device_id is passed for NCCL).
- Test execution: python test/run_test.py -> main() -> parse_args + TESTS_TO_INCLUDE env merge -> get_selected_tests (TEST groups like INDUCTOR_TESTS/DISTRIBUTED_TESTS/DYNAMO_CORE_TESTS) -> target-determination amend_tests -> do_sharding -> run_tests splits serial vs parallel files -> run_test_module invokes pytest per file with -m marker expressions; each test file uses run_tests() from common_utils which can spawn per-test-case subprocesses with a 15-min timeout and one retry.
- Codegen flow at build: cmake/Codegen.cmake runs `python -m torchgen.gen` -> main() parses native_functions.yaml + tags.yaml from aten/src/ATen -> get_custom_build_selector -> emits ATen headers/sources, dispatcher registrations, and the AOTInductor C shim into build/aten/src/ATen and torch/csrc/inductor/aoti_torch/generated.