DOCSEARCHbyAlgolia
Submit Docs

TorchAO

repository·main·Indexed 25 days ago

https://github.com/pytorch/ao

A PyTorch-native library for training-to-serving model optimization. It provides tools for sparsity and quantization (int4, int8, float8) and accelerated kernels to reduce memory usage and improve inference speed across CUDA, XPU, and CPU hardware.

Tokens
60.8K
Snippets
145
Records
258
Agent score
83%

What's inside TorchAO

  1. Overview of Blockwise FP8 Quantization

    main

    This implementation provides blockwise quantization based on the DeepSeek method. It optimizes inference by quantizing activations and weight matrices into specific block sizes using the FP8 format.

    Quantization Specifications:

    • Activations: Quantized in blocks of size 128x1.
    • Weight Matrices: Quantized in blocks of size 128x128.

    The kernels are implemented using Triton and currently support the following torch types:

    • torch.float8_e4m3fn
    • torch.float8_e5m2

  2. AWQ Quantization Overview and Benchmarks

    main

    AWQ (Activation-aware Weight Quantization) is available in torchao. The implementation includes various quantization methods such as awq-hqq-int4, awq-uint4, int4wo-hqq, and int4wo-64.

    Performance characteristics:

    • awq-hqq-int4 and int4wo-hqq utilize optimized kernels (like tinygemm int4->bf16) to provide higher throughput and lower peak memory compared to standard awq-uint4.
    • awq-uint4 does not use an efficient fused kernel, resulting in lower performance (Tokens/sec and Throughput) compared to HQQ-based methods.

    Benchmarks are provided for Llama-2-7b-chat-hf (on A100 80GB) and Llama-3-8B-Instruct (using LM eval with groupsize=128) to demonstrate the trade-offs between perplexity and inference speed.

  5. Explore TorchAO Prototype experimental kernels and utilities

    main
    The torchao.prototype module contains experimental kernels and utilities designed for efficient inference and training. These components are modularized to be composed with the torch stack and other open-source ML frameworks. Currently, the prototype includes a re-implementation of SpinQuant. Note that previously experimental low-bit optimizers have been promoted to the stable torchao.optim module.

  7. Deploy optimized models using vLLM, SGLang, and ExecuTorch

    main
    TorchAO enables an end-to-end optimization flow (pre-training, fine-tuning, and serving) by integrating quantization and sparsity techniques into partner frameworks. This tutorial series demonstrates how to use TorchAO as the underlying optimization engine for deployment via HugngingFace Transformers, vLLM, and ExecuTorch.

  8. Understand the benefits of Tensor Subclasses for quantization

    main

    In torchao, quantization is primarily implemented using tensor subclasses (classes inheriting from torch.Tensor). This approach is preferred over module swaps (replacing nn.Linear with a custom QuantizedLinear) for several reasons:

    1. Finer-grained integration: Tensor subclasses intercept computation at the function/op level (e.g., torch.add), allowing quantization to work even if the model uses functional variants or modified native modules.
    2. Better composability: Subclasses can be wrapped within one another. For example, a quantized tensor can be wrapped by a DTensor for distributed training without needing a specialized DistributedQuantizedLinear class.
    3. Reusing PyTorch components: Since quantized tensors are just torch.Tensors with different dtypes, the model structure remains unchanged, allowing existing optimization passes (like torch.compile) to work seamlessly.

  9. Use the SAM2 fork in TorchAO

    main
    TorchAO includes a specialized fork of the SAM2 model (based on commit c2ec8e14a185632b0a5d8b161928ceb50197eddc). This version includes specific modifications to support fullgraph=True during torch.compile, a mask_to_rle_pytorch_2 implementation, and performance optimizations integrated into the TorchAO packaging system.

  10. Understand CPU kernel organization in TorchAO

    main

    CPU kernels in TorchAO are organized into three distinct directories based on their target platform and implementation style:

    • torch_free_kernels: Contains kernels implemented using raw pointers, independent of PyTorch concepts like Tensor.
    • shared_kernels: Contains kernels designed to be shared between PyTorch/ATen and Executorch. These can be compiled for either platform using specific compile flags and often utilize torch_free_kernels.
    • aten_kernels: Contains kernels written specifically for PyTorch/ATen.

    When contributing new CPU kernel code, it is preferred to implement a shared kernel to maximize utility across platforms.

  11. Generate MXFP8 Roofline Performance Analysis

    main

    To generate roofline plots for Llama4 or DeepSeek V3 shapes, use the roofline_unified.py script.

    For Llama4 Shapes (K=5120, N=8192, G=8):

    cd benchmarks/prototype/moe_training/mxfp8
python roofline_unified.py --K=5120 --N=8192 --G=8 --power_limit_percent=100 --breakdown_M=131072 --plot_file=llama_rooflines.png

    For DeepSeek V3 Shapes (K=7168, N=2048, G=8):

    cd benchmarks/prototype/moe_training/mxfp8
python roofline_unified.py --K=7168 --N=2048 --G=8 --power_limit_percent=100 --breakdown_M=131072 --plot_file=dsv3_rooflines.png

  12. Serve the torchao documentation locally

    main

    After building the documentation, you can host it on a local server using sphinx-serve. This provides live reload capabilities, typically accessible at http://localhost:8000.

    sphinx-serve -b build