Implemenation of IObinding in Mixtral MoE Parity Script

onnxruntime/test/python/transformers/test_parity_mixtral_moe.py

@@ -9,10 +9,13 @@
 # Licensed under the MIT License.  See License.txt in the project root for
 # license information.
 # --------------------------------------------------------------------------
+import time
 import unittest
+import pytest
 from collections import OrderedDict
 
 import numpy
+import onnx
 import torch
 import torch.nn.functional as F
 from onnx import TensorProto, helper
@@ -23,7 +26,7 @@
 torch.manual_seed(42)
 numpy.random.seed(42)
 
-ORT_DTYPE = TensorProto.FLOAT
+ORT_DTYPE = TensorProto.FLOAT16
 NP_TYPE = numpy.float16 if ORT_DTYPE == TensorProto.FLOAT16 else numpy.float32
 THRESHOLD = 3e-2
 
@@ -80,22 +83,22 @@
             "fc1_experts_weights",
             ORT_DTYPE,
             fc1_shape,
-            fc1_experts_weights.to(torch_type).flatten().tolist(),
-            raw=False,
+            fc1_experts_weights.to(torch_type).detach().numpy().tobytes(),
+            raw=True,
         ),
         helper.make_tensor(
             "fc2_experts_weights",
             ORT_DTYPE,
             fc2_shape,
-            fc2_experts_weights.to(torch_type).flatten().tolist(),
-            raw=False,
+            fc2_experts_weights.to(torch_type).detach().numpy().tobytes(),
+            raw=True,
         ),
         helper.make_tensor(
             "fc3_experts_weights",
             ORT_DTYPE,
             fc3_shape,
-            fc3_experts_weights.to(torch_type).flatten().tolist(),
-            raw=False,
+            fc3_experts_weights.to(torch_type).detach().numpy().tobytes(),
+            raw=True,
         ),
     ]
 
@@ -126,6 +129,11 @@
     model = helper.make_model(graph)
     return model.SerializeToString()
 
+    #model_path = "mixtral_moe.onnx"
+    #onnx.save_model(model, model_path, save_as_external_data=True, all_tensors_to_one_file=True)
+
+    #return model_path
+
 
 class ClassInstantier(OrderedDict):
     def __getitem__(self, key):
@@ -216,10 +224,8 @@
         self.ffn_dim = config.intermediate_size
         self.num_experts = config.num_local_experts
         self.top_k = config.num_experts_per_tok
-
         # gating
         self.gate = nn.Linear(self.hidden_dim, self.num_experts, bias=False)
-
         self.experts = nn.ModuleList([MixtralBlockSparseTop2MLP(config) for _ in range(self.num_experts)])
 
         w1_list = []
@@ -229,11 +235,9 @@
             w1_list.append(self.experts[i].w1.weight)
             w2_list.append(self.experts[i].w2.weight)
             w3_list.append(self.experts[i].w3.weight)
-
         self.moe_experts_weight1 = torch.stack(w1_list, dim=0)
         self.moe_experts_weight2 = torch.stack(w2_list, dim=0)
         self.moe_experts_weight3 = torch.stack(w3_list, dim=0)
-
         self.batch_size = batch_size
         self.sequence_length = sequence_length
         self.moe_onnx_graph = create_moe_onnx_graph(
@@ -246,7 +250,6 @@
             self.moe_experts_weight3,
             self.top_k,
         )
-
         self.ort_sess = self.create_ort_session()
 
     def create_ort_session(self):
@@ -257,12 +260,54 @@
         cuda_providers = ["CUDAExecutionProvider"]
         if cuda_providers[0] not in onnxruntime.get_available_providers():
             return None
-
         sess_options.log_severity_level = 2
         ort_session = InferenceSession(self.moe_onnx_graph, sess_options, providers=["CUDAExecutionProvider"])
 
         return ort_session
 
+    def ort_run_with_iobinding(self, ort_inputs, repeat=1000):
+        iobinding = self.ort_sess.io_binding()
+        device_id = torch.cuda.current_device()
+
+        iobinding.bind_input(
+            name="input",
+            device_type="cuda",
+            device_id=device_id,
+            element_type=NP_TYPE,
+            shape=ort_inputs["input"].shape,
+            buffer_ptr=onnxruntime.OrtValue.ortvalue_from_numpy(ort_inputs["input"], "cuda", device_id).data_ptr(),
+        )
+
+        iobinding.bind_input(
+            name="router_probs",
+            device_type="cuda",
+            device_id=device_id,
+            element_type=NP_TYPE,
+            shape=ort_inputs["router_probs"].shape,
+            buffer_ptr=onnxruntime.OrtValue.ortvalue_from_numpy(
+                ort_inputs["router_probs"], "cuda", device_id
+            ).data_ptr(),
+        )
+
+        iobinding.bind_output(
+            name="output",
+            device_type="cuda",
+            device_id=device_id,
+            element_type=NP_TYPE,
+            shape=ort_inputs["input"].shape,
+            buffer_ptr=onnxruntime.OrtValue.ortvalue_from_numpy(
+                numpy.zeros(ort_inputs["input"].shape), "cuda", device_id
+            ).data_ptr(),
+        )
+
+        s = time.time()
+        for _ in range(repeat):
+            iobinding.synchronize_inputs()
+            self.ort_sess.run_with_iobinding(iobinding)
+            iobinding.synchronize_outputs()
+        e = time.time()
+        print(f"MoE cuda kernel time: {(e - s) / repeat * 1000} ms")
+
     def forward(self, hidden_states: torch.Tensor) -> torch.Tensor:
         """ """
         batch_size, sequence_length, hidden_dim = hidden_states.shape
@@ -305,21 +350,23 @@
         final_hidden_states = final_hidden_states.reshape(batch_size, sequence_length, hidden_dim)
         return final_hidden_states  # , router_logits
 
-    def ort_forward(self, hidden_states: torch.Tensor) -> torch.Tensor:
+    def ort_forward(self, hidden_states: torch.Tensor, iobinding=False) -> torch.Tensor:
         batch_size, sequence_length, hidden_dim = hidden_states.shape
         hidden_states = hidden_states.view(-1, hidden_dim)
         # router_logits: (batch * sequence_length, n_experts)
         router_logits = self.gate(hidden_states)
-
         ort_inputs = {
             "input": numpy.ascontiguousarray(hidden_states.detach().numpy().astype(NP_TYPE)),
             "router_probs": numpy.ascontiguousarray(router_logits.detach().numpy().astype(NP_TYPE)),
         }
 
         ort_output = None
         if self.ort_sess is not None:
-            ort_output = self.ort_sess.run(None, ort_inputs)
-            return torch.tensor(ort_output).reshape(batch_size, sequence_length, -1)  # , router_logits
+            if not iobinding:
+                ort_output = self.ort_sess.run(None, ort_inputs)
+            else:
+                self.ort_run_with_iobinding(ort_inputs)
+                return None
 
         # print_tensor("input", ort_inputs["input"])
         # print_tensor("router_probs", ort_inputs["router_probs"])
@@ -328,12 +375,12 @@
         # print_tensor("fc3_experts_weights", self.moe_experts_weight3.detach().numpy())
         # print_tensor("output", ort_output[0])
 
-        return None
+        return ort_output
 
     def parity_check(self):
         hidden_state = torch.randn(self.batch_size, self.sequence_length, self.hidden_dim)
         torch_output = self.forward(hidden_state)
-        ort_output = self.ort_forward(hidden_state)
+        ort_output = self.ort_forward(hidden_state, iobinding=True)
         if ort_output is not None:
             assert torch.allclose(torch_output, ort_output, rtol=1e-04, atol=1e-04)
             print(
@@ -346,16 +393,36 @@
                 " parity: OK",
             )
 
+    def benchmark(self):
+        self.ort_forward(iobinding=True)
+
 
 class TestMixtralMoE(unittest.TestCase):
     def test_mixtral_moe_parity(self):
-        for batch_size in [1, 16]:
-            for sequence_length in [128, 1024]:
+        for batch_size in [1, 2]:
+            for sequence_length in [8, 16]:
                 # use a small sizes to speed up the test
                 config = MixtralConfig(hidden_size=256, intermediate_size=1024)
                 mixtral_moe = MixtralSparseMoeBlock(config, batch_size, sequence_length)
                 mixtral_moe.parity_check()
 
+    @pytest.mark.slow
+    def test_mixtral_moe_large(self):
+        for batch_size in [1, 8]:
+            for sequence_length in [16, 64]:
+                config = MixtralConfig()
+                mixtral_moe = MixtralSparseMoeBlock(config, batch_size, sequence_length)
+                mixtral_moe.parity_check()
+
+    @pytest.mark.slow
+    def test_mixtral_moe_benchmark(self):
+        for batch_size in [32, 64]:
+            for sequence_length in [128, 1024]:
+                config = MixtralConfig()
+                mixtral_moe = MixtralSparseMoeBlock(config, batch_size, sequence_length)
+                mixtral_moe.benchmark()
+
+
 
 if __name__ == "__main__":
     unittest.main()