Skip to content

v0.21.0

Choose a tag to compare

View all tags
@OnlyDeniko OnlyDeniko released this 30 Jan 19:55
· 30 commits to main since this release
v0.21.0
5d60f17

RePlay 0.21.0 Release notes

Table of Contents

Release Notes

Release date: 2026-01-30 · Version: 0.21.0 · Type: Minor

Highlights

This release introduces a redesigned neural network architecture and a new data processing pipeline,
bringing improved scalability, flexibility, and transparency to model training workflow.

Key benefits:

  • Train on significantly large-scale datasets with lower memory usage thanks to batch-wise data loading without full in-memory loading
  • Customize data preprocessing more easily with composable batch-level transforms and no hidden logic
  • Build and extend models flexibly without upgrading the library using a block-based architecture with reusable components and easily adding custom blocks when needed
  • Experiment faster by decoupling blocks of models
  • Adopt new architectures incrementally, while existing pipelines continue to work

⚠️ This release introduces deprecations.
Legacy APIs continue to work, but backwards incompatible changes are planned for upcoming releases.
See the Deprecations section for details.

This release is fully backward-compatible, and existing pipelines will continue to work. It lays the foundation for future model and pipeline extensions.

Deprecations

Legacy Data Pipeline APIs

The following APIs are deprecated but continue to work in this release.
They will be removed in upcoming releases.

  • Deprecated modules

    The previous multi-stage data pipeline for neural network workflow has been deprecated, including:

    • SequentialDataset
    • SequenceTokenizer
    • SasRecTrainingDataset/SasRecValidationDataset/ SasRecPredictionDataset
    • SasRecTrainingBatch/ SasRecValidationBatch/SasRecPredictionBatch

  • Replacement

    A new data pipeline is introduced (see ParquetModule for details). The new pipeline provides greater flexibility and enables training on significantly larger datasets by avoiding out-of-memory (OOM) issues.

Model API Changes

  • SASRec APIs have been redesigned in a block-based architecture.
  • Lightning-specific wrappers are no longer model-specific.
  • Models no longer encapsulate loss computation internally.
  • Deprecated modules
    • SasRec
    • SasRecModel

Action Required

It's recommended to migrate custom datasets and preprocessing logic to the new pipeline.

New Features

ParquetModule

ParquetModule is the core building block of the new data pipeline.

  • Key Features

    • Automatic padding and sequence truncation based on a provided schema
    • Batch-wise reading and processing, enabling efficient work with large datasets in memory-constrained environments as it avoids loading the full dataset into memory.
    • Full compatibility with PyTorch Lightning Trainers.
    • Built-in support for multiple dataloaders for validation, testing and prediction
    • Built-in support for PyTorch Distributed Data Parallel (DDP)
    • Batch-level transforms can be easily composed into custom preprocessing pipelines,
      and extended with user-defined transforms when needed

  • Pipelines comparison

          LEGACY PIPELINE                            NEW PIPELINE
      ───────────────                            ────────────
         Raw Data                                  Raw Data
            ↓                                         ↓
    Filtering/Splitting                    User-defined CPU Preprocessing
            ↓                                   ├─ Filtering
         Dataset                                ├─ Splitting
            ↓                                   ├─ Tokenizing
     SequenceTokenizer                          └─ Grouping
        ├─ Tokenizing                                  ↓
        └─ Grouping                              ParquetModule
            ↓                             (GPU batch-level preprocessing)
    ├─ SasRecTrainingDataset                    ├─ Padding
        ├─ Padding                              ├─ Shifting
        ├─ Shifting                             └─ Negative Sampling
        └─ Negative Sampling                           ↓
    ├─ SasRecValidationDataset                 Model Forward Pass
    └─ SasRecPredictDataset
            ↓
  torch.utils.data.DataLoader
            ↓
     Model Forward Pass

  • Details

    Data may be prepared manually using any data processing framework, for example, Pandas, Polars, PySpark, PyArrow. It is required to save the data in Parquet format. Pay attention to partition sizes during saving the data - it is recommended to be 256 - 512 Mb.

    A ParquetModule instance is created by:

    • specifying metadata including shape and padding for each data split
    • specifying a list of batch-level transforms for each data split

    ParquetModule reads data batch-by-batch and applies split-specific transforms
    immediately before feeding the batch into the model.
    This design ensures scalability and efficient memory usage.

Neural Network Architecture Redesign

Neural network models follow a block-based architecture, where:

  • Models receive pre-built component instances instead of raw configuration parameters
  • A single unified Lightning wrapper is shared across all models
  • Core components (losses, embedders, heads, etc.) are implemented as reusable modules

All reusable building blocks are located under the replay/nn module.

The new architecture is currently implemented for SASRec and TwoTower.
Support for BERT4Rec will be added in a future release.

NN Losses and training

  • Loss computation is decoupled from NN models and is provided as standalone blocks within the block-based architecture
  • LogInCE, LogInCESampled and LogOutCE loss functions are added
  • Sampled losses (BCESampled, CESampled, LogInCESampled) support per-sample weighting via batch-provided weights
  • Support for negative labels per batch element is added for sampled losses.
  • MultiClassNegativeSamplingTransform allows negative items selection from different item catalogue subsets.

New Model: TwoTower

  • Introduced a new TwoTower model implemented using the block-based architecture
  • Fully compatible with the new data pipeline and NN workflow

Improvements

  • Optimization of SequenceEncodingRule. The achieved acceleration on large datasets (100+ million rows) in an industrial pipeline is up to 15+ times.

Bug Fixes

  • Fixed SasRecCompiled and Bert4RecCompiled compilation issues with torch >= 2.9.0

Migration Notes

  • Migration to the new parquet-based pipeline will be required for existing datasets and preprocessing pipelines in upcoming releases
  • Custom preprocessing logic from SasRecTrainingDataset and related classes
    should be reimplemented using batch-level transforms
  • Existing models (for example, replay.models.nn.sequential.SasRec) that were trained in earlier versions of the library are deprecated but still available.
    However, checkpoints produced by the legacy SasRec cannot be loaded into a redesigned SasRec.
  • For the redesigned SasRec and new TwoTower, the functionality for compiling models in ONNX format and further compilation via OpenVINO is not yet available. We will implement this functionality in the next release.
    But for now, you can use the to_onnx() method from LightningModule or apply the torch.onnx.export() method to the model itself.
  • NN Transformers (SasRec and BERT4Rec) now use dict-based batches instead of NamedTuple

References

For detailed usage examples and documentation, see the links.

Assets 2
Loading