Skip to content

Latest commit

 

History

History
110 lines (92 loc) · 4.79 KB

README.md

File metadata and controls

110 lines (92 loc) · 4.79 KB
Raw

About Model_To_MLIR

This project was developed during my summer internship at the NYUAD University under the Supervision of Dr. Riyadh Baghdadi

Project Overview

  • Project Name: Model_To_MLIR
  • Python Version: 3.10.12

Setup Instructions

  1. Create Virtual Environment: 
    python -m venv myvenv
  2. Activate The Virtual Environment 
    source ./myvenv/bin/activate
  3. Install Required Packages: 
    pip install -r requirements.txt
  4. Install Torch-MLIR: 
    pip install --pre torch-mlir torchvision \
--extra-index-url https://download.pytorch.org/whl/nightly/cpu \
-f https://github.com/llvm/torch-mlir-release/releases/expanded_assets/dev-wheels

Using The repository :

Check ./examples folder:

  1. In this project i used the FX importer since all the other importers are disabled for the time being

Explanation of ./examples/bert.py example

    import torch
    from torch_mlir import fx
    from transformers import BertForMaskedLM
    from contextlib import redirect_stdout


    # Define a wrapper class for the BERT model
    class BertWrapper(torch.nn.Module):
        def __init__(self):
            super().__init__()
            self.model = BertForMaskedLM.from_pretrained("bert-base-uncased", return_dict=False)

        def forward(self, input_ids):
            return self.model(input_ids)

    # Instantiate and prepare the model
    model = BertWrapper()
    model.eval()

    # Create dummy input simulating tokenized text input for BERT
    # The input IDs should be in the range of the vocabulary size (0-30522 for BERT)
    dummy_input = torch.randint(0, 30522, (1, 10))  # Adjust size as needed

    # Export to MLIR using fx.export_and_import
    module = fx.export_and_import(
        model,
        dummy_input,
        output_type="linalg-on-tensors",  # 
        func_name=model.__class__.__name__,
    )

    # Save the exported MLIR module to a file
    out_mlir_path = "./bert_linalg.mlir"
    with open(out_mlir_path, 'w') as f:
        with redirect_stdout(f):
            print(module.operation.get_asm())

This example demonstrates how to convert a BERT model to MLIR (Multi-Level Intermediate Representation) using the torch_mlir library. The script performs the following steps:

  1. Import necessary libraries:

    • torch: PyTorch library for tensor computations and deep learning.
    • torch_mlir: Library for exporting PyTorch models to MLIR.
    • transformers: Hugging Face library for transformer models.
    • contextlib.redirect_stdout: Utility to redirect standard output.

  2. Define a wrapper class for the BERT model:

    • BertWrapper: A custom PyTorch module that wraps the BertForMaskedLM model from the transformers library.

  3. Instantiate and prepare the model:

    • Create an instance of BertWrapper.
    • Set the model to evaluation mode using model.eval().

  4. Create dummy input:

    • Generate a tensor simulating tokenized text input for BERT. The input IDs are randomly generated within the vocabulary size range (0-30522).

  5. Export to MLIR:

    • Use fx.export_and_import to export the model to MLIR format with the specified output type (linalg-on-tensors) and function name.
    • There are many other dialects that are supported like torch, tosa, stablehlo, raw

  6. Save the exported MLIR module to a file:

    • Write the MLIR representation to a file named bert_linalg.mlir using redirect_stdout to capture the output of module.operation.get_asm().

Remark:

Depending on the model size, the generated MLIR file size will vary. For NLP models, it can take a lot of space. For LLMs, this repo does not support them yet as it requires significant GPU and CPU resources.

Resources

  1. For more in-depth understanding of the concept, check the following readme file: touch-mlir architecture
  2. Refer to the torch-mlir Discord channel in the LLVM project Discord server, you will find a lot of helpful people over there.

Limitation Of the project:

This project is intended to help you do the conversion from PyTorch to MLIR. Although it covers a significant portion of the conversion process, it is not perfect. You need to test the generated code using the mlir-opt command and fix any syntax errors by yourself.

Example of Usage

mlir-opt ./linear_regression_linalg.mlir