DOCSEARCHbyAlgolia
Submit Docs

TorchRec Documentation

repository·main·Indexed 25 days ago

https://github.com/meta-pytorch/torchrec

A library for recommendation systems featuring EmbeddingBagCollection (EBC) and FusedEmbeddingBagCollection (Fused EBC) for high-performance embedding lookups. It includes support for DistributedModelParallel (DMP) and Data Parallel (DDP) modes, UVM-caching for large tables, and a dynamic embedding extension (torchrec_dynamic_embedding) for connecting to external Parameter Servers. The repository also provides a BERT4Rec implementation example demonstrating transformer-based sequential recommendation.

Tokens
75.4K
Snippets
134
Records
299
Agent score
82%

What's inside TorchRec

  1. Overview of TorchRec Inference Library

    main
    TorchRec Inference is an experimental C++ library designed for multi-gpu inference. It works by sharding models packaged in Python using torch.package. To serve inference from C++ while subverting the Global Interpreter Lock (GIL), it utilizes the torch.deploy library, which launches multiple Python interpreters to carry the packaged model.

  2. Overview of TorchRec features

    main

    TorchRec is a PyTorch domain library designed for large-scale recommender systems (RecSys). It provides primitives for handling large embedding tables sharded across multiple GPUs/nodes.

    Key features include:

    • Parallelism Primitives: Enables hybrid data-parallelism and model-parallelism for multi-device/multi-node models.
    • Sharders: Supports various embedding table sharding strategies: data-parallel, table-wise, row-wise, table-wise-row-wise, column-wise, and table-wise-column-wise.
    • Planner: Automatically generates optimized sharding plans.
    • Pipelined Training: Overlaps dataloading (device transfer), inter-device communications (input_dist), and computation (forward/backward) to increase performance.
    • Optimized Kernels: Powered by FBGEMM.
    • Quantization & Inference: Supports reduced precision training/inference and optimization for C++ inference.
    • RecSys Modules & Datasets: Includes common modules and datasets like Criteo click logs and MovieLens.

  3. Overview of the TorchRec Bert4Rec Example

    main

    The bert4rec_main.py script provides an implementation of BERT4REC using TorchRec. This example demonstrates a model that combines both data parallel components (such as transformation blocks) and model parallel components (such as item embeddings).

    Key implementation details:

    • Datasets: The model can be run using a random dataloader or the MovieLens dataset.
    • Output Layer: Unlike the original paper which uses the embedding layer weights for the final output via matmul, this implementation uses a linear layer as the final output layer to achieve higher performance.

  4. Overview of TorchRec Inference

    main
    TorchRec Inference is a C++ library designed for multi-GPU inference. It works by sharding models that were originally written and packaged in Python using torch.package. To serve these models from C++ while subverting the Global Interpreter Lock (GIL), it utilizes the torch.deploy library, which launches multiple Python interpreters to carry the packaged models.

  5. Overview of TorchRec capabilities

    main

    TorchRec is a PyTorch domain library designed for building and scaling recommendation systems (RecSys). It provides specialized building blocks for handling industry-scale workloads, specifically focusing on sparse features and large-scale model distribution.

    Key capabilities include:

    • Embedding modules: Optimized high-performance embedding bags and tables for sparse categorical features.
    • Sharding and parallelism: Native support for model, table, and row-wise sharding across multiple GPUs and nodes.
    • Distributed training: Deep integration with PyTorch Distributed and FSDP (Fully Sharded Data Parallel) for training massive models.
    • Feature processing: Utilities for managing RecSys-specific data structures, including jagged tensors and pooling operations.

  6. Understand the TorchRec DLRM Golden Training Example

    main

    The golden_training example demonstrates production-ready training patterns for DLRM (Deep Learning Recommendation Model) using TorchRec.

    Key TorchRec Features Demonstrated:

    • EmbeddingBagCollection: Efficient embedding lookups for multiple categorical features.
    • DistributedModelParallel: Automatic sharding of embedding tables across multiple GPUs.
    • TrainPipelineSparseDist: Overlaps communication, computation, and data transfer to hide latency.
    • RowWiseAdagrad: A row-wise optimizer for sparse embeddings, fused with the backward pass.
    • QCommsConfig: Quantized communication (FP16/BF16) to optimize multi-node network bandwidth.

    Production Patterns:

    • SPMD Training: Each process runs the same script with different data shards.
    • Model Parallelism: Sharding large embedding tables across GPUs.
    • Pipeline Parallelism: Overlapping data loading, communication, and computation.
    • Mixed Precision: Using FP16 for forward and BF16 for backward communication.

  7. Bert4Rec Example Directory Structure

    main

    The Bert4Rec example repository is organized as follows:

    • bert4rec_main.py: The main training script.
    • bert4rec_metrics.py: Contains evaluation metrics such as HR@K and NDCG@K.
    • data/: Contains bert4rec_movielens_datasets.py for MovieLens data loading.
    • dataloader/: Contains bert4rec_movielens_dataloader.py for batch preparation with masking.
    • models/: Contains bert4rec.py for the model definition.
    • tests/: Contains unit tests for the main script.

  8. TorchRec Overview

    main

    TorchRec is a PyTorch-based library specifically designed for building, scaling, and deploying large-scale recommendation system models. It provides specialized primitives that address the unique challenges of recommendation systems, such as massive embedding tables and high-scale distributed training, which are not the primary focus of standard PyTorch.

    Key capabilities include:

    • Specialized Components: Modules optimized for recommendation tasks, particularly embedding tables.
    • Advanced Sharding: Flexible methods for sharding embedding tables (e.g., Row-Wise, Column-Wise, Table-Wise) with automatic device topology optimization.
    • Distributed Training: Sophisticated model parallelism techniques tailored for massive recommendation scales.
    • Optimized Performance: Training and inference components are highly optimized using FBGEMM.
    • Deployment Path: APIs for transforming trained models for inference and loading them into C++ environments for optimal performance.
    • PyTorch Integration: Seamlessly integrates with the existing PyTorch ecosystem, tools, and workflows.

  9. Key TorchRec features used in Two-Tower Retrieval

    main

    The Two-Tower retrieval example demonstrates several core TorchRec capabilities:

    • EmbeddingBagCollection: Efficient embedding lookups for categorical features.
    • DistributedModelParallel: Model-parallel sharding of embedding tables across multiple GPUs.
    • TrainPipelineSparseDist: Overlapped communication and computation to accelerate training.
    • KeyedJaggedTensor: An efficient sparse tensor format designed for variable-length features.
    • RowWiseAdagrad: A specialized row-wise optimizer for efficient embedding training.
    • torchrec.quant: Support for INT8 quantization to enable efficient inference.

  10. Transfer learning with pretrained embeddings in TorchRec

    main

    This example demonstrates how to train a distributed recommendation model using TorchRec by leveraging pretrained embeddings instead of training from scratch.

    Key Concepts

    • Initialization: Instead of random noise, embeddings are initialized with meaningful pretrained values (e.g., Word2Vec, GloVe, Item2Vec, or embeddings from related tasks).
    • Efficiency: For large pretrained embedding tables, the example utilizes the share_memory_ API. This allows multiple processes to efficiently access the same tensors from shared memory, reducing memory overhead during the loading phase.
    • Goal: The objective is to achieve faster convergence with fewer training epochs by starting from a meaningful representation space.

  11. What is Multi-Probe Zero Collision Hash (MPZCH)?

    main

    Multi-probe Zero Collision Hash (MPZCH) is a technique used in TorchRec to reduce collision rates during embedding table lookups. It manages hash collisions by using two essential tables:

    1. Identity Table: Maps an input hash value to a remapped ID. The value stored in each slot is the input hash value, and the remapped ID is the index of that slot.
    2. Metadata Table: Shares the same length as the identity table and records the timestamp when a hash value was inserted into its corresponding identity table slot.

    The MPZCH lookup/insertion process involves two probes:

    • First Probe: Checks for available or evictable slots in the identity table.
    • Second Probe: Checks if the slot indexed by the input hash value is occupied. If empty, the hash is inserted. If occupied, it performs a linear probe to find the next available slot. If all slots are full, it finds the next evictable slot (based on a time threshold) and replaces the expired hash value.

  12. What is a JaggedTensor?

    main

    A JaggedTensor is a specialized data type in TorchRec designed to represent sparse features with variable-length sequences efficiently. Unlike a standard torch.Tensor, which requires padding to make all sequences the same length, a JaggedTensor stores data contiguously without padding, saving memory and computation.

    It consists of three key components:

    • Lengths: A list of integers representing the number of elements for each entity.
    • Offsets: A list of integers representing the starting index of each sequence in the flattened values tensor (an alternative to Lengths).
    • Values: A 1D tensor containing the actual values for each entity, stored contiguously.
    # User interactions:
# - User 1 interacted with 2 items
# - User 2 interacted with 3 items
# - User 3 interacted with 1 item
lengths = [2, 3, 1]
offsets = [0, 2, 5]  # Starting index of each user's interactions
values = torch.Tensor([101, 102, 201, 202, 203, 301])  # Item IDs interacted with
jt = JaggedTensor(lengths=lengths, values=values)
# OR
jt = JaggedTensor(offsets=offsets, values=values)