crystalai/ai-auto-train-deep-learning-multi-dimensional-multi-model-create-Transformational-tools-app

1# SPDX-FileCopyrightText: Copyright (c) 2022-2024 NVIDIA CORPORATION & AFFILIATES. All rights reserved.2# SPDX-License-Identifier: Apache-2.03#4# Licensed under the Apache License, Version 2.0 (the "License");5# you may not use this file except in compliance with the License.6# You may obtain a copy of the License at7#8# http://www.apache.org/licenses/LICENSE-2.09#10# Unless required by applicable law or agreed to in writing, software11# distributed under the License is distributed on an "AS IS" BASIS,12# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.13# See the License for the specific language governing permissions and14# limitations under the License.15import unittest16 17import numpy as np18import torch19from polygraphy.backend.trt import EngineFromNetwork, TrtRunner20from torch import nn21 22import tensorrt_llm23from tensorrt_llm import Module, Tensor24 25 26class TorchMLP(nn.Module):27 28    def __init__(self, hidden_size, ffn_hidden_size, bias=True):29        super().__init__()30        self.fc = nn.Linear(hidden_size, ffn_hidden_size, bias=bias)31        self.proj = nn.Linear(ffn_hidden_size, hidden_size, bias=bias)32 33    def forward(self, hidden_states):34        inter = self.fc(hidden_states)35        inter = nn.functional.relu(inter)36        output = self.proj(inter)37        return output, inter38 39 40class MLP(Module):41 42    def __init__(self,43                 hidden_size,44                 ffn_hidden_size,45                 bias=True,46                 tp_group=None,47                 tp_size=1):48        super().__init__()49        self.fc = tensorrt_llm.layers.ColumnLinear(hidden_size,50                                                   ffn_hidden_size,51                                                   bias=bias,52                                                   tp_group=tp_group,53                                                   tp_size=tp_size,54                                                   gather_output=False)55        self.proj = tensorrt_llm.layers.RowLinear(ffn_hidden_size,56                                                  hidden_size,57                                                  bias=bias,58                                                  tp_group=tp_group,59                                                  tp_size=tp_size)60 61    def forward(self, hidden_states):62        inter = self.fc(hidden_states)63        inter = tensorrt_llm.functional.relu(inter)64        self.register_network_output('inter', inter)65        output = self.proj(inter)66        return output67 68 69class TestDebuggingAPI(unittest.TestCase):70 71    def setUp(self):72        tensorrt_llm.logger.set_level('error')73 74    def test_debugging_api(self):75        # test data76        dtype = 'float32'77        hidden_size = 76878        x_data = torch.randn(2, 16, hidden_size)79 80        tm = TorchMLP(hidden_size=hidden_size,81                      ffn_hidden_size=hidden_size * 4,82                      bias=False)83 84        # construct trt network85        builder = tensorrt_llm.Builder()86        net = builder.create_network()87        with tensorrt_llm.net_guard(net):88            x = Tensor(name='x',89                       shape=x_data.shape,90                       dtype=tensorrt_llm.str_dtype_to_trt(dtype))91 92            gm = MLP(hidden_size=hidden_size,93                     ffn_hidden_size=4 * hidden_size,94                     bias=False)95            gm.fc.weight.value = tm.fc.weight.detach().cpu().numpy()96            gm.proj.weight.value = tm.proj.weight.detach().cpu().numpy()97 98            output = gm.forward(x)99            net._mark_output(output, 'output',100                             tensorrt_llm.str_dtype_to_trt(dtype))101 102            for k, v in gm.named_network_outputs():103                net._mark_output(v, k, tensorrt_llm.str_dtype_to_trt(dtype))104 105        # trt run106        build_engine = EngineFromNetwork((builder.trt_builder, net.trt_network))107        with TrtRunner(build_engine) as runner:108            outputs = runner.infer(feed_dict={'x': x_data.numpy()})109 110        # pytorch run111        with torch.no_grad():112            ref1, ref2 = tm(x_data)113 114        # compare diff115        np.testing.assert_allclose(ref1.cpu().numpy(),116                                   outputs['output'],117                                   atol=1e-5)118        np.testing.assert_allclose(ref2.cpu().numpy(),119                                   outputs['inter'],120                                   atol=1e-5)121