DOCSEARCHbyAlgolia
Submit Docs

PyTorch

repository·develop·Indexed 42 days ago

https://github.com/rocm/pytorch

A Python-based high-performance tensor library and deep learning framework featuring strong GPU acceleration and a dynamic, tape-based autograd system. This documentation covers Android installation via Gradle (Lite Interpreter and Full JIT builds), building from source, linking to native C++ builds, the c10 dispatcher architecture, and registering custom operators.

Tokens
360.4K
Snippets
919
Records
1.4K
Agent score
53%

What's inside PyTorch

  1. Overview of the PyTorch C++ Frontend

    develop

    The PyTorch C++ frontend is a C++17 library designed for CPU and GPU tensor computation. It provides automatic differentiation and high-level building blocks for machine learning, mirroring the functionality of the Python frontend. It is intended for use cases where the Python interpreter is unsuitable, such as low-latency, high-performance, or highly multithreaded production environments (e.g., video games or production servers).

    The frontend consists of several core components:

    ComponentDescription
    torch::TensorAutomatically differentiable, efficient CPU and GPU enabled tensors
    torch::nnA collection of composable modules for neural network modeling
    torch::optimOptimization algorithms like SGD, Adam or RMSprop to train your models
    torch::dataDatasets, data pipelines and multi-threaded, asynchronous data loader
    torch::serializeA serialization API for storing and loading model checkpoints
    torch::pythonGlue to bind your C++ models into Python
    torch::jitPure C++ access to the TorchScript JIT compiler

  2. What is the NumPy <> PyTorch Compat Layer?

    develop

    The torch._numpy module is a standalone implementation of the NumPy public API using PyTorch tensors. It does not depend on the NumPy library itself.

    Primary Use Case: It is used by TorchDynamo to trace through NumPy code and lower it into PyTorch code, allowing NumPy operations to be part of a compiled PyTorch graph.

    Warning: This layer is not intended as a general-purpose replacement for NumPy or the standard PyTorch API. When used in eager mode (without torch.compile), it is significantly slower than regular NumPy or PyTorch.

  3. Overview of torch.distributed.pipelining

    develop

    torch.distributed.pipelining is a toolkit for implementing Pipeline Parallelism (PP). It allows partitioning a model's execution across multiple devices so that different micro-batches can execute different parts of the model concurrently. This is particularly useful for large-scale training, bandwidth-limited clusters, and large model inference where computation cannot hide communication overhead.

    The package consists of two main components:

    1. Splitting Frontend: Takes model code and automatically (or manually) splits it into "model partitions" while capturing data-flow relationships.
    2. Distributed Runtime: Executes pipeline stages on different devices, handling micro-batch splitting, scheduling, communication, and gradient propagation.

    Note: This package is currently in alpha state and was migrated from the PiPPy project. APIs are subject to change.

  4. Overview of the PyTorch C++ API components

    develop

    The PyTorch C++ API is composed of five primary building blocks that allow for tensor computation and dynamic neural network development in C++:

    1. ATen: The foundational tensor and mathematical operation library. It provides the core Tensor class and hundreds of operations.
    2. Autograd: Extends ATen with automatic differentiation capabilities by recording operations to form an autograd graph.
    3. C++ Frontend: A high-level modeling interface (similar to the Python API) for defining neural networks, optimizers, and data pipelines.
    4. TorchScript: An interface to the TorchScript JIT compiler and interpreter, allowing for model serialization and execution in C++ environments without Python.
    5. C++ Extensions: A mechanism to extend the Python API with custom C++ or CUDA routines, providing integration with Python setuptools and JIT compilation.

  5. What is Static Runtime and when to use it

    develop

    Static Runtime is an optimized CPU inference runtime for PyTorch models designed as a drop-in replacement for the TorchScript JIT interpreter in both C++ and Python.

    Warning: This is an experimental feature.

    It is most effective when:

    1. The model has minimal control flow.
    2. PyTorch overhead (like tensor creation and allocation) is a significant portion of runtime. Static Runtime coalesces intermediate tensor memory into a single slab to avoid dynamic allocations.
    3. Inference performance is a critical requirement.

    Current Assumptions:

    • Execution is CPU-only.
    • Inference-only execution.
    • Static input dtypes.
    • Static input shapes (while dynamic shapes are supported, they may degrade performance).

  6. Introduction to PyTorch DTensor

    develop

    PyTorch DTensor (DistributedTensor) is a prototype implementation of distributed tensor primitives designed for the SPMD (Single Program Multiple Devices) paradigm. It allows users to express complex tensor distributions using sharding and replication strategies, enabling native Tensor Parallelism and seamless interoperability between different parallelism strategies like Data Parallel (DDP/FSDP) and Pipeline Parallelism.

    Key value propositions include:

    • Uniform Checkpointing: A consistent way to save/load state_dict regardless of complex sharding/replication strategies.
    • Eager Mode Tensor Parallelism: Provides flexibility to mix sharding and replication in eager mode.
    • SPMD Foundation: Serves as a building block for compiler-based distributed training.
    # Example: Sharding a big tensor across devices
import os
import torch
from torch.distributed.tensor import init_device_mesh, Shard, distribute_tensor

# Create a mesh topology with the available devices
mesh = init_device_mesh("cuda", (int(os.environ["WORLD_SIZE"]),))
big_tensor = torch.randn(100000, 88)

# Shard this tensor over the mesh by sharding `big_tensor`'s 0th dimension over the 0th dimension of `mesh`.
my_dtensor = distribute_tensor(big_tensor, mesh, [Shard(dim=0)])

  7. What is QNNPACK?

    develop

    QNNPACK (Quantized Neural Networks PACKage) is a mobile-optimized library designed for low-precision, high-performance neural network inference. It provides implementations of common neural network operators specifically for quantized 8-bit tensors.

    Note that QNNPACK is intended as a provider of low-level performance primitives for high-level deep learning frameworks (like PyTorch) rather than a direct tool for machine learning researchers.

  8. Overview of `torch.compile()` Benchmarking suites

    develop

    The benchmarks/dynamo/ directory contains benchmarking code for TorchDynamo and various backends (including TorchInductor). It is organized into three main benchmark suites:

    1. TorchBenchmark: A diverse set of models seeded from research models. Requires installation of the pytorch/benchmark repository. Use the provided Makefile to set up the environment to match PyTorch CI versions.
    2. HuggingFace Models: Primarily transformer models. The runner huggingface.py automatically handles dependency installation on the first run.
    3. TIMM Models: Primarily vision models. The runner timm_models.py automatically handles dependency installation on the first run.

  9. Overview of the Distributed RPC Framework

    develop

    The Distributed RPC (Remote Procedure Call) framework enables multi-machine model training by providing primitives for remote communication and a high-level API for automatic differentiation across machines.

    It is composed of four primary API sets:

    1. Remote Procedure Call (RPC): Executes functions on remote workers (rpc_sync, rpc_async, or remote).
    2. Remote Reference (RRef): Acts as a distributed shared pointer to objects on remote workers, handling reference counting transparently.
    3. Distributed Autograd: Automatically manages the backward pass across multiple workers by stitching together local autograd engines.
    4. Distributed Optimizer: Wraps local optimizers on different workers into a single interface to update parameters scattered across the cluster.

    Important Limitations:

    • CUDA Support is Beta: Many features, including RRefs, JIT compatibility, distributed autograd, and distributed optimizer, are currently unsupported or discouraged when using CUDA tensors.
    • Tensor Transmission: When passing Tensor objects, the framework transmits metadata (shape, stride, etc.). Direct transmission of CUDA tensors is disallowed to prevent crashes due to device mismatch; users should move tensors to CPU before transmission and to the target device on the callee side.

  10. Overview of Sparse benchmarks

    develop
    The sparse benchmarks are designed to evaluate the performance of sparse matrix functionality within PyTorch. They are primarily used to compare the performance of sparse matrix routines, such as SpMV (Sparse Matrix-Vector multiplication), across different sparse matrix formats and against other frameworks like TensorFlow.

  11. Introduction to PyTorch Symmetric Memory

    develop

    PyTorch Symmetric Memory (torch.distributed._symmetric_memory) is an alpha-state feature designed to allow developers to program modern GPU systems (connected via high-bandwidth links like NVLink, InfiniBand, or RoCE) as a single, unified memory space.

    It unlocks three primary capabilities:

    1. Customized communication patterns: Allows writing custom kernels for specific algorithms and data types.
    2. In-kernel compute-comm fusion: Enables device-initiated communication, allowing computation and data movement to be fused at a fine granularity.
    3. Low-latency remote access: Leverages protocols like RDMA to enable direct, high-speed communication between processors and memory, bypassing traditional CPU/network stack overhead.

    Note: The API is currently under development and subject to change.

  12. Overview of CUTLASS extensions in PyTorch

    develop

    The aten/src/ATen/native/cuda/cutlass_extensions/ directory contains a subset of files from the FasterTransformer project. These extensions are used to support mixed datatype GEMM (General Matrix Multiply) implementations within the aten/src/ATen/native/cuda/MixedDTypesLinear.cu file of the PyTorch source tree.

    Key details:

    • Purpose: Provides specialized kernels for mixed-precision operations.
    • Origin: Based on commit f8e42aa of FasterTransformer, modified to support CUTLASS 3.x (the original project used CUTLASS 2.10).
    • Maintenance Note: These files are linted according to PyTorch standards and may differ from the original FasterTransformer source. The goal is to eventually remove this directory once CUTLASS officially incorporates this functionality.