Repository: alibaba/BladeDISC (language: python)

Issue:
add bert training example 
Add a bert training example with TorchBlade accelerator in the codebase and check the loss on CI. This task needs to check the disc backend with LTC tracing modular at least, and Dynamo in the future.

cc @tanyokwok 


Discussion / hints:

Fix patch:
diff --git a/examples/PyTorch/Training/bert_train_ltc.py b/examples/PyTorch/Training/bert_train_ltc.py
new file mode 100644
index 00000000000..fc7e3f9bb0e
--- /dev/null
+++ b/examples/PyTorch/Training/bert_train_ltc.py
@@ -0,0 +1,281 @@
+# Copyright 2022 The BladeDISC Authors. All rights reserved.
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+# http://www.apache.org/licenses/LICENSE-2.0
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+import argparse
+import inspect
+import os
+import sys
+import time
+from datetime import timedelta
+import timeit
+import pandas as pd
+import numpy as np
+
+import torch
+from datasets import load_dataset
+from datasets import load_metric
+from torch.utils.data import DataLoader
+from transformers import AutoModelForSequenceClassification
+from transformers import AutoTokenizer
+
+import torch_blade
+import torch_blade.utils as utils
+import ctypes
+import torch._lazy as ltc
+import torch._lazy.metrics as ltc_metrics
+utils.disable_pytorch_jit()
+torch_blade.init_ltc_disc_backend()
+torch.manual_seed(2)
+
+
+torch.backends.cuda.matmul.allow_tf32 = True
+
+# You will download around 84G dataset if you run this end to end training/evaluation example.
+
+os.environ["TOKENIZERS_PARALLELISM"] = "false"
+
+results = []
+def printStats(backend, timings, precision):
+    times = np.array(timings)
+    steps = len(times)
+    speeds = 1.0/times
+    time_mean = np.mean(times)
+    time_med = np.median(times)
+    time_99th = np.percentile(times, 99)
+    time_std = np.std(times, ddof=0)
+    speed_mean = np.mean(speeds)
+    speed_med = np.median(speeds)
+
+    msg = (
+        "\n%s =================================\n"
+        "num iterations=%d\n"
+        "  Median FPS: %.1f, mean: %.1f\n"
+        "  Median latency: %.6f, mean: %.6f, 99th_p: %.6f, std_dev: %.6f\n"
+    ) % (
+        backend,
+        steps,
+        speed_med,
+        speed_mean,
+        time_med,
+        time_mean,
+        time_99th,
+        time_std,
+    )
+    meas = {
+        "Backend": backend,
+        "precision": precision,
+        "Median(FPS)": speed_med,
+        "Mean(FPS)": speed_mean,
+        "Median(ms)": time_med,
+        "Mean(ms)": time_mean,
+        "99th_p": time_99th,
+        "std_dev": time_std,
+    }
+    results.append(meas)
+
+
+def data_processing(num_samples, batch_size):
+    dataset = load_dataset("yelp_review_full")
+    tokenizer = AutoTokenizer.from_pretrained("bert-base-cased")
+
+    def tokenize_function(examples):
+        return tokenizer(examples["text"], padding="max_length", truncation=True)
+
+    tokenized_datasets = dataset.map(tokenize_function, batched=True)
+
+    tokenized_datasets = tokenized_datasets.remove_columns(["text"])
+    tokenized_datasets = tokenized_datasets.rename_column("label", "labels")
+    tokenized_datasets.set_format("torch")
+
+    small_train_dataset = tokenized_datasets["train"].select(range(num_samples))
+    small_eval_dataset = tokenized_datasets["test"].select(range(num_samples))
+
+    train_dataloader = DataLoader(small_train_dataset, batch_size=batch_size)
+    eval_dataloader = DataLoader(small_eval_dataset, batch_size=batch_size)
+
+    return train_dataloader, eval_dataloader
+
+
+def training_iter_fn(batch, model, optimizer, scaler):
+    with torch.cuda.amp.autocast():
+        outputs = model(**batch)
+        loss = outputs.loss
+    scaler.scale(loss).backward()
+    return loss
+
+
+def get_device(backend):
+    # reference to Lazy device if enable LazyTensor Core
+    device = torch.device("cuda") if torch.cuda.is_available() else torch.device("cpu")
+    if backend == "ltc-disc":
+        device = torch.device("lazy")
+    return device
+
+
+def model_training_evaluation(
+        backend, train_dataloader, eval_dataloader, model, optimizer, num_epochs, evaluation):
+    device = get_device(backend)
+
+    model.to(device)
+    model.train()
+    total_iter = 0
+    timings = []
+    profiling_warmup = 20
+
+    tr_loss = 0.0
+    loss_history = []
+    for epoch in range(num_epochs):
+        end_time = timeit.default_timer()
+        for i, batch in enumerate(train_dataloader, 0):
+            total_iter += 1
+            start_time = end_time
+            batch = {k: v.to(device) for k, v in batch.items()}
+            optimizer.zero_grad()
+            loss = model(**batch).loss
+            loss.backward()
+            if backend == "ltc-disc":
+                torch._lazy.mark_step()
+            loss_history.append(loss.item())
+            tr_loss += loss.item()
+            if i % 5 == 0:
+                tr_loss = torch.tensor(0.0).to(device)
+
+            end_time = timeit.default_timer()
+            meas_time = end_time - start_time
+            print("epoch: {}, batch: {}, loss: {}, elapsed:{}".format(epoch, i, loss.item(), meas_time))
+            timings.append(meas_time)
+
+    printStats(backend or "PyTorch", timings[1:], "fp32")
+    if evaluation:
+        metric = load_metric("accuracy")
+        model.eval()
+        opt_model = model
+        for batch in eval_dataloader:
+            batch = {k: v.to(device) for k, v in batch.items()}
+            with torch.no_grad():
+                outputs = opt_model(**batch)
+
+            logits = outputs.logits
+            predictions = torch.argmax(logits, dim=-1)
+            metric.add_batch(predictions=predictions, references=batch["labels"])
+
+        return loss_history, metric.compute()
+    else:
+        return loss_history, None
+
+
+def check_loss(ref_loss, res_loss):
+    assert len(ref_loss) == len(res_loss)
+    length = len(ref_loss)
+    x = min(length, 10)
+    if sum(res_loss[-x:]) / 10 <= sum(ref_loss[-x:]) / 10 + 1e-1:
+        return True
+    else:
+        return False
+
+
+def parse_args():
+    parser = argparse.ArgumentParser(
+        description="TorchDynamo end to end training/evaluation benchmark"
+    )
+    parser.add_argument(
+        "--epochs", type=int, default=10, help="number of epochs to train (default: 10)"
+    )
+    parser.add_argument(
+        "--num-samples",
+        type=int,
+        default=1000,
+        help="number of samples to train/eval (default: 1000)",
+    )
+    parser.add_argument(
+        "--batch-size",
+        type=int,
+        default=48,
+        help="input batch size for training (default: 8)",
+    )
+    parser.add_argument(
+        "--lr", type=float, default=5e-5, help="learning rate (default: 5e-5)"
+    )
+    parser.add_argument(
+        "--acc-backend",
+        default="eager",
+        help="train model with accelerator backend[eager, ltc-disc] (default: eager)",
+    )
+
+    parser.add_argument(
+        "--optimizer",
+        default="Adam",
+        help="train model using a given optimizer (default: Adam)",
+    )
+    parser.add_argument(
+        "--evaluation",
+        action="store_true",
+        help="running evaluation after model training",
+    )
+    args = parser.parse_args()
+    return args
+
+
+def main():
+    args = parse_args()
+    train_dataloader, eval_dataloader = data_processing(
+        args.num_samples, args.batch_size
+    )
+    model = AutoModelForSequenceClassification.from_pretrained(
+        "bert-base-cased", num_labels=5
+    )
+    optimizer_cls = getattr(sys.modules["torch.optim"], args.optimizer)
+    if "capturable" in inspect.signature(optimizer_cls).parameters.keys():
+        optimizer = optimizer_cls(model.parameters(), lr=args.lr, capturable=True)
+    else:
+        optimizer = optimizer_cls(model.parameters(), lr=args.lr)
+    if args.acc_backend == "ltc-disc":
+        actual_loss, accuracy = model_training_evaluation(
+            args.acc_backend,
+            train_dataloader,
+            eval_dataloader,
+            model,
+            optimizer,
+            args.epochs,
+            args.evaluation
+        )
+
+    native_start = time.time()
+    expect_loss, accuracy = model_training_evaluation(
+        None,
+        train_dataloader,
+        eval_dataloader,
+        model,
+        optimizer,
+        args.epochs,
+        args.evaluation
+    )
+    native_end = time.time()
+
+    check_success = check_loss(actual_loss, expect_loss)
+    if args.evaluation:
+        print(f"Model accuracy: {accuracy}")
+    native_elapsed = native_end - native_start
+    print(
+        f"Train model on {args.epochs} epochs with backend {args.acc_backend} and optimizer {args.optimizer}:"
+    )
+    print(f"PyTorch spent {timedelta(seconds=native_elapsed/args.epochs)} per epoch")
+
+    # Generate report
+    print("Model Summary:")
+    summary = pd.DataFrame(results)
+    print(summary.to_markdown())
+
+    assert check_success, "PyTorch LTC end to end training loss is greater than native Pytorch"
+
+
+if __name__ == "__main__":
+    main()
diff --git a/examples/PyTorch/Training/requirements.txt b/examples/PyTorch/Training/requirements.txt
new file mode 100644
index 00000000000..9e2ffddd003
--- /dev/null
+++ b/examples/PyTorch/Training/requirements.txt
@@ -0,0 +1,4 @@
+transformers==4.24.0
+pandas==1.5.1
+tabulate==0.9.0
+datasets==2.6.1
\ No newline at end of file
diff --git a/pytorch_blade/pytorch_blade/compiler/jit/torch/shape_analysis.cpp b/pytorch_blade/pytorch_blade/compiler/jit/torch/shape_analysis.cpp
index 4aed6da73de..4df3c1e152b 100644
--- a/pytorch_blade/pytorch_blade/compiler/jit/torch/shape_analysis.cpp
+++ b/pytorch_blade/pytorch_blade/compiler/jit/torch/shape_analysis.cpp
@@ -2112,7 +2112,8 @@ class ShapePropagator : public PropertyPropBase {
     } else if (
         node->matches(
             "aten::embedding_dense_backward(Tensor grad_output, Tensor indices, int num_weights, int padding_idx, bool scale_grad_by_freq) -> Tensor")) {
-      if (auto type = node->input(0)->type()->cast<TensorType>()) {
+      auto type = node->input(0)->type()->cast<TensorType>();
+      if (type && type->dim()) {
         auto numWeightsOptional = node->get<int64_t>(attr::num_weights);
         if (type->isComplete()) {
           std::vector<int64_t> newSizes = {numWeightsOptional.value()};
diff --git a/pytorch_blade/pytorch_blade/compiler/mlir/converters/mhlo_conversion.cpp b/pytorch_blade/pytorch_blade/compiler/mlir/converters/mhlo_conversion.cpp
index 10636827328..cc8876735df 100644
--- a/pytorch_blade/pytorch_blade/compiler/mlir/converters/mhlo_conversion.cpp
+++ b/pytorch_blade/pytorch_blade/compiler/mlir/converters/mhlo_conversion.cpp
@@ -180,6 +180,7 @@ ConvertTorchToMhlo(std::shared_ptr<torch::jit::Graph> graph) {
 
   mlir::OpPrintingFlags print_flags;
   print_flags.elideLargeElementsAttrs();
+  print_flags.enableDebugInfo();
 
   bool enable_printing =
       env::ReadBoolFromEnvVar("TORCH_BLADE_MHLO_DEBUG_LOG", false);
diff --git a/pytorch_blade/pytorch_blade/compiler/mlir/converters/torch_mlir_op_filter.cpp b/pytorch_blade/pytorch_blade/compiler/mlir/converters/torch_mlir_op_filter.cpp
index 142c82ff4e2..9ce253ac54b 100644
--- a/pytorch_blade/pytorch_blade/compiler/mlir/converters/torch_mlir_op_filter.cpp
+++ b/pytorch_blade/pytorch_blade/compiler/mlir/converters/torch_mlir_op_filter.cpp
@@ -40,6 +40,7 @@ const std::unordered_set<std::string> &GetTorchMlirWhiteList() {
   static std::unordered_set<std::string> white_list{
       "aten::_autocast_to_reduced_precision",
       "aten::__and__",
+      "aten::_softmax",
       "aten::add",
       "aten::addmm",
       "aten::arange",
@@ -91,7 +92,8 @@ const std::unordered_set<std::string> &GetTorchMlirWhiteList() {
       "aten::ne",
       "aten::neg",
       "aten::neg",
-      "aten::native_dropout",
+      // TODO(disc): need to lower mhlo::RngOp to mhlo_disc::UniformOp
+      //"aten::native_dropout",
       "aten::permute",
       "aten::pow",
       "aten::relu",
diff --git a/pytorch_blade/pytorch_blade/torch-mlir/lib/Conversion/TorchToMhlo/DiscTorchToMhlo.cpp b/pytorch_blade/pytorch_blade/torch-mlir/lib/Conversion/TorchToMhlo/DiscTorchToMhlo.cpp
index f122d99796a..012925b1072 100644
--- a/pytorch_blade/pytorch_blade/torch-mlir/lib/Conversion/TorchToMhlo/DiscTorchToMhlo.cpp
+++ b/pytorch_blade/pytorch_blade/torch-mlir/lib/Conversion/TorchToMhlo/DiscTorchToMhlo.cpp
@@ -1102,7 +1102,6 @@ class DiscConvertTorchToMhlo
     INSERT_UNARY_PATTERN(AtenNegOp, mhlo::NegOp)
     INSERT_UNARY_PATTERN(AtenFloorOp, mhlo::FloorOp)
     INSERT_UNARY_PATTERN(AtenCeilOp, mhlo::CeilOp)
-    INSERT_UNARY_PATTERN(AtenItemOp, tensor::ExtractOp)
 #undef INSERT_UNARY_PATTERN
 
 #define INSERT_EXTRACT_PATTERN(AtenOp) \
@@ -1110,6 +1109,7 @@ class DiscConvertTorchToMhlo
   patterns.add<ConvertAtenExtractOp<AtenOp>>(typeConverter, context);
     INSERT_EXTRACT_PATTERN(AtenScalarImplicitOp)
     INSERT_EXTRACT_PATTERN(AtenIntTensorOp)
+    INSERT_EXTRACT_PATTERN(AtenItemOp)
 #undef INSERT_EXTRACT_PATTERN
 
 #define INSERT_ATENOP_PATTERN(AtenOp) \
diff --git a/pytorch_blade/scripts/build_pytorch_blade.sh b/pytorch_blade/scripts/build_pytorch_blade.sh
index 1e1b830572a..08afaa56d30 100755
--- a/pytorch_blade/scripts/build_pytorch_blade.sh
+++ b/pytorch_blade/scripts/build_pytorch_blade.sh
@@ -62,5 +62,20 @@ function ci_build() {
     python3 setup.py bdist_wheel;
 }
 
+function test_example() {
+  if [ "$TORCH_BLADE_CI_BUILD_TORCH_VERSION" == "1.12.0+cu113" ]; then
+    # note: BladeDISC only tested LTC-DISC backend on PyTorch 1.12
+    # There is a but that causing OOM on LTC, so BladeDISC only check minor iterations,
+    # issue: https://github.com/pytorch/pytorch/issues/80942
+    (
+      cd ../examples/PyTorch/Training && \
+      python -m pip install -r requirements.txt && \
+      LTC_DISC_CUDA=1 python bert_train_ltc.py --batch-size 8 --acc-backend ltc-disc --num-samples 152 --epochs 1 2>&1
+    )
+  fi
+}
+
 # Build
 ci_build
+
+test_example


Test patch:
diff --git a/pytorch_blade/tests/ltc/test_mnist.py b/pytorch_blade/tests/ltc/test_mnist.py
index 79f8b749602..3dd55a607f0 100644
--- a/pytorch_blade/tests/ltc/test_mnist.py
+++ b/pytorch_blade/tests/ltc/test_mnist.py
@@ -25,6 +25,7 @@
 if torch_blade._is_ltc_available:
     import torch._lazy as ltc
     torch_blade.init_ltc_disc_backend()
+
 import unittest
 
 LOG_INTERVAL=100