已合并
[feat]support atb_llm whl包生成 #205
zhaokerui创建于 1月10日
[feat]support atb_llm whl包生成 #205
已合并
共 48 个文件变更+112-87
The file is empty
The file is empty
The file is empty
The file is empty
The file is empty
The file is empty
The file is empty
The file is empty
The file is empty
| @@ -7,6 +7,7 @@ | |||
| 7 | # EITHER EXPRESS OR IMPLIED, INCLUDING BUT NOT LIMITED TO NON-INFRINGEMENT, | 7 | # EITHER EXPRESS OR IMPLIED, INCLUDING BUT NOT LIMITED TO NON-INFRINGEMENT, |
| 8 | # MERCHANTABILITY OR FIT FOR A PARTICULAR PURPOSE. | 8 | # MERCHANTABILITY OR FIT FOR A PARTICULAR PURPOSE. |
| 9 | # See the Mulan PSL v2 for more details. | 9 | # See the Mulan PSL v2 for more details. |
| 10 | + | ||
| 10 | from typing import Optional | 11 | from typing import Optional |
| 11 | from enum import Enum | 12 | from enum import Enum |
| 12 | import importlib | 13 | import importlib |
The file is empty
The file is empty
The file is empty
The file is empty
The file is empty
The file is empty
The file is empty
The file is empty
The file is empty
The file is empty
The file is empty
The file is empty
The file is empty
The file is empty
The file is empty
The file is empty
The file is empty
The file is empty
The file is empty
The file is empty
The file is empty
The file is empty
The file is empty
The file is empty
The file is empty
The file is empty
The file is empty
The file is empty
The file is empty
The file is empty
The file is empty
The file is empty
The file is empty
The file is empty
The file is empty
| @@ -1,5 +1,5 @@ | |||
| 1 | #!/bin/bash | 1 | #!/bin/bash |
| 2 | -# Copyright (c) Huawei Technologies Co., Ltd. 2023. All rights reserved. | 2 | +# Copyright (c) Huawei Technologies Co., Ltd. 2023-2026. All rights reserved. |
| 3 | # | 3 | # |
| 4 | # Licensed under the Apache License, Version 2.0 (the "License"); | 4 | # Licensed under the Apache License, Version 2.0 (the "License"); |
| 5 | # you may not use this file except in compliance with the License. | 5 | # you may not use this file except in compliance with the License. |
| @@ -359,65 +359,6 @@ EOF | |||
| 359 | 359 | ||
| 360 | } | 360 | } |
| 361 | 361 | ||
| 362 | -function fn_build_for_ci() | ||
| 363 | -{ | ||
| 364 | - fn_files_authority_limit | ||
| 365 | - cd $OUTPUT_DIR/atb_models | ||
| 366 | - rm -rf ./*.tar.gz | ||
| 367 | - cp $ATB_MODELS_DIR/dist/atb_llm*.whl . | ||
| 368 | - cp -r $ATB_MODELS_DIR/atb_llm . | ||
| 369 | - cp -r $ATB_MODELS_DIR/docs . | ||
| 370 | - cp $ATB_MODELS_DIR/setup.py . | ||
| 371 | - cp -r $ATB_MODELS_DIR/examples . | ||
| 372 | - cp -r $ATB_MODELS_DIR/tests . | ||
| 373 | - cp -r $ATB_MODELS_DIR/requirements . | ||
| 374 | - cp -r $ATB_MODELS_DIR/public_address_statement.md . | ||
| 375 | - cp $README_DIR/README.md . | ||
| 376 | - fn_build_version_info | ||
| 377 | - | ||
| 378 | - torch_vision=$(pip list | grep torch | head -n 1 | awk '{print $2}' | cut -d '+' -f1) | ||
| 379 | - if [ "$USE_CXX11_ABI" == "OFF" ];then | ||
| 380 | - abi=0 | ||
| 381 | - else | ||
| 382 | - abi=1 | ||
| 383 | - fi | ||
| 384 | - | ||
| 385 | - TMP_VERSION=$(python3 -c 'import sys; print(sys.version_info[0], ".", sys.version_info[1])' | tr -d ' ') | ||
| 386 | - PY_MINOR_VERSION=${TMP_VERSION##*.} | ||
| 387 | - PY_VERSION="py3$PY_MINOR_VERSION" | ||
| 388 | - | ||
| 389 | - tar_package_name="Ascend-mindie-atb-models_${PACKAGE_NAME}_linux-${ARCH}_${PY_VERSION}_torch${torch_vision}-abi${abi}.tar.gz" | ||
| 390 | - | ||
| 391 | - if [ $IS_RELEASE -eq 1 ]; then | ||
| 392 | - source_folder_list=$(cat $SCRIPT_DIR/release_folder.ini | xargs) | ||
| 393 | - chmod 750 $source_folder_list | ||
| 394 | - find $source_folder_list -mindepth 1 -type d -exec chmod 550 {} \; | ||
| 395 | - find $source_folder_list -type f \( -name "*.py" -o -name "*.sh" -o -name "*.so" -o -name "*.tar.gz" \) -exec chmod 550 {} \; | ||
| 396 | - find $source_folder_list -type f \( -name "*.json" -o -name "*.jsonl" -o -name "*.csv" -o -name "*.txt" \) -exec chmod 640 {} \; | ||
| 397 | - tar czpf $tar_package_name $source_folder_list --owner=0 --group=0 | ||
| 398 | - else | ||
| 399 | - tar czpf $tar_package_name ./* --owner=0 --group=0 | ||
| 400 | - fi | ||
| 401 | - | ||
| 402 | - if [ -f "README.md" ];then | ||
| 403 | - rm -rf README.md | ||
| 404 | - fi | ||
| 405 | - | ||
| 406 | - if [ $IS_RELEASE -eq 1 ]; then | ||
| 407 | - cd $OUTPUT_DIR | ||
| 408 | - mkdir -p debug_symbols | ||
| 409 | - debug_symbols_package_name="$OUTPUT_DIR/debug_symbols/Ascend-mindie-atb-models-debug-symbols_${PACKAGE_NAME}_linux-${ARCH}_${PY_VERSION}_torch${torch_vision}-abi${abi}.tar.gz" | ||
| 410 | - tar czpf $debug_symbols_package_name atb_models_debug_symbols | ||
| 411 | - echo "Save debug symbols file to $OUTPUT_DIR/debug_symbols" | ||
| 412 | - fi | ||
| 413 | -} | ||
| 414 | - | ||
| 415 | -function fn_make_whl() { | ||
| 416 | - echo "make atb_llm whl package" | ||
| 417 | - cd $ATB_MODELS_DIR | ||
| 418 | - python3 $ATB_MODELS_DIR/setup.py bdist_wheel | ||
| 419 | -} | ||
| 420 | - | ||
| 421 | function fn_extract_debug_symbols() { | 362 | function fn_extract_debug_symbols() { |
| 422 | local in_dir=$1 | 363 | local in_dir=$1 |
| 423 | local out_dir=$2 | 364 | local out_dir=$2 |
| @@ -488,18 +429,19 @@ print(pybind11_cmake_dir) | |||
| 488 | if [ "$CLEAN_FIRST" == "ON" ];then | 429 | if [ "$CLEAN_FIRST" == "ON" ];then |
| 489 | make clean | 430 | make clean |
| 490 | fi | 431 | fi |
| 432 | + # Determine parallel build level | ||
| 433 | + PARALLEL_LEVEL=${CMAKE_BUILD_PARALLEL_LEVEL:-8} | ||
| 434 | + | ||
| 491 | if [ "$USE_VERBOSE" == "ON" ];then | 435 | if [ "$USE_VERBOSE" == "ON" ];then |
| 492 | - VERBOSE=1 make -j | 436 | + VERBOSE=1 make -j${PARALLEL_LEVEL} |
| 493 | else | 437 | else |
| 494 | - make -j | 438 | + make -j${PARALLEL_LEVEL} |
| 495 | fi | 439 | fi |
| 496 | make install | 440 | make install |
| 497 | - fn_make_whl | ||
| 498 | 441 | ||
| 499 | if [ $IS_RELEASE -eq 1 ]; then | 442 | if [ $IS_RELEASE -eq 1 ]; then |
| 500 | fn_extract_debug_symbols $OUTPUT_DIR/atb_models $OUTPUT_DIR/atb_models_debug_symbols/lib | 443 | fn_extract_debug_symbols $OUTPUT_DIR/atb_models $OUTPUT_DIR/atb_models_debug_symbols/lib |
| 501 | fi | 444 | fi |
| 502 | - fn_build_for_ci | ||
 | |||
| 503 | } | 445 | } |
| 504 | 446 | ||
| 505 | function fn_main() | 447 | function fn_main() |
| @@ -1,4 +1,4 @@ | |||
| 1 | -# Copyright (c) Huawei Technologies Co., Ltd. 2024-2025. All rights reserved. | 1 | +# Copyright (c) Huawei Technologies Co., Ltd. 2024-2026. All rights reserved. |
| 2 | # MindIE is licensed under Mulan PSL v2. | 2 | # MindIE is licensed under Mulan PSL v2. |
| 3 | # You can use this software according to the terms and conditions of the Mulan PSL v2. | 3 | # You can use this software according to the terms and conditions of the Mulan PSL v2. |
| 4 | # You may obtain a copy of Mulan PSL v2 at: | 4 | # You may obtain a copy of Mulan PSL v2 at: |
| @@ -7,20 +7,102 @@ | |||
| 7 | # EITHER EXPRESS OR IMPLIED, INCLUDING BUT NOT LIMITED TO NON-INFRINGEMENT, | 7 | # EITHER EXPRESS OR IMPLIED, INCLUDING BUT NOT LIMITED TO NON-INFRINGEMENT, |
| 8 | # MERCHANTABILITY OR FIT FOR A PARTICULAR PURPOSE. | 8 | # MERCHANTABILITY OR FIT FOR A PARTICULAR PURPOSE. |
| 9 | # See the Mulan PSL v2 for more details. | 9 | # See the Mulan PSL v2 for more details. |
| 10 | -from setuptools import setup | ||
| 11 | 10 | ||
| 11 | +import os | ||
| 12 | +import subprocess | ||
| 13 | +import logging | ||
| 14 | +import shutil | ||
| 15 | +from pathlib import Path | ||
| 16 | +from setuptools.command.build_py import build_py as _build_py | ||
| 17 | +from setuptools import setup, find_packages | ||
| 18 | +from wheel.bdist_wheel import bdist_wheel as _bdist_wheel | ||
| 19 | + | ||
| 20 | +logging.basicConfig(level=logging.INFO) | ||
| 21 | +os.environ["SOURCE_DATE_EPOCH"] = "315532800" # 315532800 means 1980-01-01 00:00:00 UTC | ||
| 22 | + | ||
| 23 | + | ||
| 24 | +class CustomBuildPy(_build_py): | ||
| 25 | + def run(self): | ||
| 26 | + project_root = Path(__file__).resolve().parent | ||
| 27 | + build_dir = project_root / "scripts" | ||
| 28 | + | ||
| 29 | + logging.info(">>> Running build.sh to compile shared libraries...") | ||
| 30 | + subprocess.check_call(["bash", "build.sh"], cwd=str(build_dir)) | ||
| 31 | + | ||
| 32 | + build_pkg = Path(self.build_lib) / "atb_llm" | ||
| 33 | + (build_pkg / "lib").mkdir(parents=True, exist_ok=True) | ||
| 34 | + shutil.copytree("output/atb_models/lib", build_pkg / "lib", dirs_exist_ok=True) | ||
| 35 | + | ||
| 36 | + (build_pkg / "conf").mkdir(parents=True, exist_ok=True) | ||
| 37 | + shutil.copytree("atb_llm/conf", build_pkg / "conf", dirs_exist_ok=True) | ||
| 38 | + extra_dirs = ["docs", "requirements", "tools", "examples", "tests"] | ||
| 39 | + for d in extra_dirs: | ||
| 40 | + src = project_root / d | ||
| 41 | + if src.exists(): | ||
| 42 | + dst = build_pkg / d | ||
| 43 | + logging.info(f"Copying {src} -> {dst}") | ||
| 44 | + shutil.copytree(src, dst, dirs_exist_ok=True) | ||
| 45 | + else: | ||
| 46 | + logging.warning(f"Source directory {src} does not exist, skipping.") | ||
| 47 | + super().run() | ||
| 48 | + | ||
| 49 | + | ||
| 50 | +class BDistWheel(_bdist_wheel): | ||
| 51 | + def finalize_options(self): | ||
| 52 | + super().finalize_options() | ||
| 53 | + self.root_is_pure = False | ||
| 54 | + | ||
| 55 | + | ||
| 56 | +def get_version() -> str: | ||
| 57 | + """ | ||
| 58 | + Return version string. | ||
| 59 | + | ||
| 60 | + Priority: | ||
| 61 | + 1. Environment variable MINDIE_LLM_VERSION_OVERRIDE | ||
| 62 | + 2. Default version | ||
| 63 | + """ | ||
| 64 | + version = os.getenv("MINDIE_LLM_VERSION_OVERRIDE", "1.0.0") | ||
| 65 | + logging.info(f"Use mindie llm version: {version}") | ||
| 66 | + return version | ||
| 67 | + | ||
| 68 | + | ||
| 69 | +def use_cxx11_abi() -> str: | ||
| 70 | + """ | ||
| 71 | + Return whether to use CXX11 ABI as a string ("0" or "1"). | ||
| 72 | + Uses torch.compiled_with_cxx11_abi() if available. | ||
| 73 | + Defaults to "0" if torch is not found or detection fails. | ||
| 74 | + """ | ||
| 75 | + try: | ||
| 76 | + import torch # noqa: F401 | ||
| 77 | + abi = torch.compiled_with_cxx11_abi() | ||
| 78 | + if abi is not None: | ||
| 79 | + abi = str(int(bool(abi))) | ||
| 80 | + logging.info(f"Detect ABI from torch, set USE_CXX11_ABI to {abi}") | ||
| 81 | + return abi | ||
| 82 | + except Exception: | ||
| 83 | + logging.warning("No torch detected on current environment.") | ||
| 84 | + return "0" | ||
| 12 | 85 | ||
| 13 | setup( | 86 | setup( |
| 14 | name="atb_llm", | 87 | name="atb_llm", |
| 15 | - version="0.0.1", | 88 | + version=get_version(), |
| 16 | author="", | 89 | author="", |
| 17 | author_email="", | 90 | author_email="", |
| 18 | - description="ATB LLM Project", | 91 | + description="Atb models Project", |
| 19 | long_description="", | 92 | long_description="", |
| 20 | - package_dir={'atb_llm': 'atb_llm'}, | 93 | + install_requires=[], |
| 21 | - package_data={ | ||
| 22 | - '': ['*.xlsx', '*.h5', '*.csv', '*.so', '*.avsc', '*.xml', '*.pkl', '*.sql', '*.ini', '*.json'] | ||
| 23 | - }, | ||
| 24 | zip_safe=False, | 94 | zip_safe=False, |
| 25 | - python_requires=">=3.7", | 95 | + python_requires=">=3.10", |
| 26 | -) | 96 | + include_package_data=False, |
| 97 | + packages=find_packages(), | ||
| 98 | + package_data={ | ||
| 99 | + "atb_llm": [ | ||
| 100 | + "lib/*" | ||
| 101 | + ] | ||
| 102 | + }, | ||
| 103 | + cmdclass={ | ||
| 104 | + "build_py": CustomBuildPy, | ||
| 105 | + "bdist_wheel": BDistWheel | ||
| 106 | + }, | ||
| 107 | + has_ext_modules=lambda: True | ||
| 108 | +) | ||
| @@ -93,7 +93,7 @@ class TestLoraLayers(unittest.TestCase): | |||
| 93 | mock_weight_loader_func.return_value = linear_tensor | 93 | mock_weight_loader_func.return_value = linear_tensor |
| 94 | linear_layer = FakeParallelLinear(["linear"]) | 94 | linear_layer = FakeParallelLinear(["linear"]) |
| 95 | lora_layer = ColumnParallelLinearWithLoRA(linear_layer) | 95 | lora_layer = ColumnParallelLinearWithLoRA(linear_layer) |
| 96 | - self.assertTrue(torch.equal(lora_layer.weight, linear_tensor)) | 96 | + self.assertTrue(torch.allclose(lora_layer.weight, linear_tensor)) |
| 97 | self.assertIsNone(lora_layer.bias) | 97 | self.assertIsNone(lora_layer.bias) |
| 98 | 98 | ||
| 99 | 99 | ||
| @@ -119,8 +119,8 @@ class TestLoraLayers(unittest.TestCase): | |||
| 119 | lora_b = torch.zeros(self.max_loras + 1, dim_r, self.n, dtype=self.dtype) | 119 | lora_b = torch.zeros(self.max_loras + 1, dim_r, self.n, dtype=self.dtype) |
| 120 | self.assertIsInstance(lora_layer.lora_a_stacked, Parameter) | 120 | self.assertIsInstance(lora_layer.lora_a_stacked, Parameter) |
| 121 | self.assertIsInstance(lora_layer.lora_b_stacked, Parameter) | 121 | self.assertIsInstance(lora_layer.lora_b_stacked, Parameter) |
| 122 | - self.assertTrue(torch.equal(lora_layer.lora_a_stacked.cpu(), lora_a)) | 122 | + self.assertTrue(torch.allclose(lora_layer.lora_a_stacked.cpu(), lora_a)) |
| 123 | - self.assertTrue(torch.equal(lora_layer.lora_b_stacked.cpu(), lora_b)) | 123 | + self.assertTrue(torch.allclose(lora_layer.lora_b_stacked.cpu(), lora_b)) |
| 124 | 124 | ||
| 125 | 125 | ||
| 126 | def test_parallel_linear_with_lora_set_lora(self, mock_weight_loader_func): | 126 | def test_parallel_linear_with_lora_set_lora(self, mock_weight_loader_func): |
| @@ -137,8 +137,8 @@ class TestLoraLayers(unittest.TestCase): | |||
| 137 | lora_b = torch.rand((self.r, self.n), device=self.device, dtype=torch.float16) | 137 | lora_b = torch.rand((self.r, self.n), device=self.device, dtype=torch.float16) |
| 138 | index = random.randint(0, self.max_loras - 1) | 138 | index = random.randint(0, self.max_loras - 1) |
| 139 | lora_layer.set_lora(index, lora_a, lora_b) | 139 | lora_layer.set_lora(index, lora_a, lora_b) |
| 140 | - self.assertTrue(torch.equal(lora_layer.lora_a_stacked[index, :self.r].cpu(), lora_a.cpu())) | 140 | + self.assertTrue(torch.allclose(lora_layer.lora_a_stacked[index, :self.r].cpu(), lora_a.cpu())) |
| 141 | - self.assertTrue(torch.equal(lora_layer.lora_b_stacked[index, :self.r].cpu(), lora_b.cpu())) | 141 | + self.assertTrue(torch.allclose(lora_layer.lora_b_stacked[index, :self.r].cpu(), lora_b.cpu())) |
| 142 | 142 | ||
| 143 | 143 | ||
| 144 | def test_parallel_linear_with_lora_reset_lora(self, mock_weight_loader_func): | 144 | def test_parallel_linear_with_lora_reset_lora(self, mock_weight_loader_func): |
| @@ -155,8 +155,8 @@ class TestLoraLayers(unittest.TestCase): | |||
| 155 | lora_layer.reset_lora(index) | 155 | lora_layer.reset_lora(index) |
| 156 | lora_a = torch.zeros(dim_r, self.k, dtype=self.dtype) | 156 | lora_a = torch.zeros(dim_r, self.k, dtype=self.dtype) |
| 157 | lora_b = torch.zeros(dim_r, self.n, dtype=self.dtype) | 157 | lora_b = torch.zeros(dim_r, self.n, dtype=self.dtype) |
| 158 | - self.assertTrue(torch.equal(lora_layer.lora_a_stacked[index].cpu(), lora_a)) | 158 | + self.assertTrue(torch.allclose(lora_layer.lora_a_stacked[index].cpu(), lora_a)) |
| 159 | - self.assertTrue(torch.equal(lora_layer.lora_b_stacked[index].cpu(), lora_b)) | 159 | + self.assertTrue(torch.allclose(lora_layer.lora_b_stacked[index].cpu(), lora_b)) |
| 160 | 160 | ||
| 161 | 161 | ||
| 162 | 162 | ||
| @@ -216,8 +216,8 @@ class TestLoraLayers(unittest.TestCase): | |||
| 216 | lora_layer.dtype = self.dtype | 216 | lora_layer.dtype = self.dtype |
| 217 | lora_a = lora_layer.load_lora_a(mock_weight_tool_obj, ["lora_A"]) | 217 | lora_a = lora_layer.load_lora_a(mock_weight_tool_obj, ["lora_A"]) |
| 218 | lora_b = lora_layer.load_lora_b(mock_weight_tool_obj, ["lora_B"], [1]) | 218 | lora_b = lora_layer.load_lora_b(mock_weight_tool_obj, ["lora_B"], [1]) |
| 219 | - self.assertTrue(torch.equal(lora_a, tensor_2)) | 219 | + self.assertTrue(torch.allclose(lora_a, tensor_2)) |
| 220 | - self.assertTrue(torch.equal(lora_b, tensor_1.T.contiguous())) | 220 | + self.assertTrue(torch.allclose(lora_b, tensor_1.T.contiguous())) |
| 221 | mock_weight_tool_obj.get_sharded.assert_called_once() | 221 | mock_weight_tool_obj.get_sharded.assert_called_once() |
| 222 | mock_weight_tool_obj.get_tensor.assert_called_once() | 222 | mock_weight_tool_obj.get_tensor.assert_called_once() |
| 223 | 223 | ||
| @@ -235,8 +235,8 @@ class TestLoraLayers(unittest.TestCase): | |||
| 235 | lora_layer.dtype = self.dtype | 235 | lora_layer.dtype = self.dtype |
| 236 | lora_a = lora_layer.load_lora_a(mock_weight_tool_obj, ["lora1_A", "lora2_A"]) | 236 | lora_a = lora_layer.load_lora_a(mock_weight_tool_obj, ["lora1_A", "lora2_A"]) |
| 237 | lora_b = lora_layer.load_lora_b(mock_weight_tool_obj, ["lora1_B", "lora2_B"], [1, 1]) | 237 | lora_b = lora_layer.load_lora_b(mock_weight_tool_obj, ["lora1_B", "lora2_B"], [1, 1]) |
| 238 | - self.assertTrue(torch.equal(lora_a, torch.cat([tensor_2] * 2, dim=0))) | 238 | + self.assertTrue(torch.allclose(lora_a, torch.cat([tensor_2] * 2, dim=0))) |
| 239 | - self.assertTrue(torch.equal(lora_b, torch.block_diag(*([tensor_1] * 2)).T.contiguous())) | 239 | + self.assertTrue(torch.allclose(lora_b, torch.block_diag(*([tensor_1] * 2)).T.contiguous())) |
| 240 | 240 | ||
| 241 | def test_row_parallel_linear_with_lora(self): | 241 | def test_row_parallel_linear_with_lora(self): |
| 242 | linear_tensor = torch.rand((self.n, self.k), device=self.device, dtype=torch.float16) | 242 | linear_tensor = torch.rand((self.n, self.k), device=self.device, dtype=torch.float16) |
| @@ -252,8 +252,8 @@ class TestLoraLayers(unittest.TestCase): | |||
| 252 | lora_layer.dtype = self.dtype | 252 | lora_layer.dtype = self.dtype |
| 253 | lora_a = lora_layer.load_lora_a(mock_weight_tool_obj, ["lora_A"]) | 253 | lora_a = lora_layer.load_lora_a(mock_weight_tool_obj, ["lora_A"]) |
| 254 | lora_b = lora_layer.load_lora_b(mock_weight_tool_obj, ["lora_B"], [1]) | 254 | lora_b = lora_layer.load_lora_b(mock_weight_tool_obj, ["lora_B"], [1]) |
| 255 | - self.assertTrue(torch.equal(lora_a, tensor_1)) | 255 | + self.assertTrue(torch.allclose(lora_a, tensor_1)) |
| 256 | - self.assertTrue(torch.equal(lora_b, tensor_2.T.contiguous())) | 256 | + self.assertTrue(torch.allclose(lora_b, tensor_2.T.contiguous())) |
| 257 | mock_weight_tool_obj.get_sharded.assert_called_once() | 257 | mock_weight_tool_obj.get_sharded.assert_called_once() |
| 258 | mock_weight_tool_obj.get_tensor.assert_called_once() | 258 | mock_weight_tool_obj.get_tensor.assert_called_once() |
| 259 | 259 | ||
【review】函数调用点删除后，原来的函数就没有用了，考虑是否同步删除一下。