diff --git a/PyTorch/built-in/foundation/CodeGeeX2/LICENSE b/PyTorch/built-in/foundation/CodeGeeX2/LICENSE new file mode 100644 index 0000000000000000000000000000000000000000..ac4aee55c0698300d21541d5395d452016585f7a --- /dev/null +++ b/PyTorch/built-in/foundation/CodeGeeX2/LICENSE @@ -0,0 +1,201 @@ + Apache License + Version 2.0, January 2004 + http://www.apache.org/licenses/ + + TERMS AND CONDITIONS FOR USE, REPRODUCTION, AND DISTRIBUTION + + 1. Definitions. + + "License" shall mean the terms and conditions for use, reproduction, + and distribution as defined by Sections 1 through 9 of this document. + + "Licensor" shall mean the copyright owner or entity authorized by + the copyright owner that is granting the License. + + "Legal Entity" shall mean the union of the acting entity and all + other entities that control, are controlled by, or are under common + control with that entity. For the purposes of this definition, + "control" means (i) the power, direct or indirect, to cause the + direction or management of such entity, whether by contract or + otherwise, or (ii) ownership of fifty percent (50%) or more of the + outstanding shares, or (iii) beneficial ownership of such entity. + + "You" (or "Your") shall mean an individual or Legal Entity + exercising permissions granted by this License. + + "Source" form shall mean the preferred form for making modifications, + including but not limited to software source code, documentation + source, and configuration files. + + "Object" form shall mean any form resulting from mechanical + transformation or translation of a Source form, including but + not limited to compiled object code, generated documentation, + and conversions to other media types. + + "Work" shall mean the work of authorship, whether in Source or + Object form, made available under the License, as indicated by a + copyright notice that is included in or attached to the work + (an example is provided in the Appendix below). + + "Derivative Works" shall mean any work, whether in Source or Object + form, that is based on (or derived from) the Work and for which the + editorial revisions, annotations, elaborations, or other modifications + represent, as a whole, an original work of authorship. For the purposes + of this License, Derivative Works shall not include works that remain + separable from, or merely link (or bind by name) to the interfaces of, + the Work and Derivative Works thereof. + + "Contribution" shall mean any work of authorship, including + the original version of the Work and any modifications or additions + to that Work or Derivative Works thereof, that is intentionally + submitted to Licensor for inclusion in the Work by the copyright owner + or by an individual or Legal Entity authorized to submit on behalf of + the copyright owner. For the purposes of this definition, "submitted" + means any form of electronic, verbal, or written communication sent + to the Licensor or its representatives, including but not limited to + communication on electronic mailing lists, source code control systems, + and issue tracking systems that are managed by, or on behalf of, the + Licensor for the purpose of discussing and improving the Work, but + excluding communication that is conspicuously marked or otherwise + designated in writing by the copyright owner as "Not a Contribution." + + "Contributor" shall mean Licensor and any individual or Legal Entity + on behalf of whom a Contribution has been received by Licensor and + subsequently incorporated within the Work. + + 2. Grant of Copyright License. Subject to the terms and conditions of + this License, each Contributor hereby grants to You a perpetual, + worldwide, non-exclusive, no-charge, royalty-free, irrevocable + copyright license to reproduce, prepare Derivative Works of, + publicly display, publicly perform, sublicense, and distribute the + Work and such Derivative Works in Source or Object form. + + 3. Grant of Patent License. Subject to the terms and conditions of + this License, each Contributor hereby grants to You a perpetual, + worldwide, non-exclusive, no-charge, royalty-free, irrevocable + (except as stated in this section) patent license to make, have made, + use, offer to sell, sell, import, and otherwise transfer the Work, + where such license applies only to those patent claims licensable + by such Contributor that are necessarily infringed by their + Contribution(s) alone or by combination of their Contribution(s) + with the Work to which such Contribution(s) was submitted. If You + institute patent litigation against any entity (including a + cross-claim or counterclaim in a lawsuit) alleging that the Work + or a Contribution incorporated within the Work constitutes direct + or contributory patent infringement, then any patent licenses + granted to You under this License for that Work shall terminate + as of the date such litigation is filed. + + 4. Redistribution. You may reproduce and distribute copies of the + Work or Derivative Works thereof in any medium, with or without + modifications, and in Source or Object form, provided that You + meet the following conditions: + + (a) You must give any other recipients of the Work or + Derivative Works a copy of this License; and + + (b) You must cause any modified files to carry prominent notices + stating that You changed the files; and + + (c) You must retain, in the Source form of any Derivative Works + that You distribute, all copyright, patent, trademark, and + attribution notices from the Source form of the Work, + excluding those notices that do not pertain to any part of + the Derivative Works; and + + (d) If the Work includes a "NOTICE" text file as part of its + distribution, then any Derivative Works that You distribute must + include a readable copy of the attribution notices contained + within such NOTICE file, excluding those notices that do not + pertain to any part of the Derivative Works, in at least one + of the following places: within a NOTICE text file distributed + as part of the Derivative Works; within the Source form or + documentation, if provided along with the Derivative Works; or, + within a display generated by the Derivative Works, if and + wherever such third-party notices normally appear. The contents + of the NOTICE file are for informational purposes only and + do not modify the License. You may add Your own attribution + notices within Derivative Works that You distribute, alongside + or as an addendum to the NOTICE text from the Work, provided + that such additional attribution notices cannot be construed + as modifying the License. + + You may add Your own copyright statement to Your modifications and + may provide additional or different license terms and conditions + for use, reproduction, or distribution of Your modifications, or + for any such Derivative Works as a whole, provided Your use, + reproduction, and distribution of the Work otherwise complies with + the conditions stated in this License. + + 5. Submission of Contributions. Unless You explicitly state otherwise, + any Contribution intentionally submitted for inclusion in the Work + by You to the Licensor shall be under the terms and conditions of + this License, without any additional terms or conditions. + Notwithstanding the above, nothing herein shall supersede or modify + the terms of any separate license agreement you may have executed + with Licensor regarding such Contributions. + + 6. Trademarks. This License does not grant permission to use the trade + names, trademarks, service marks, or product names of the Licensor, + except as required for reasonable and customary use in describing the + origin of the Work and reproducing the content of the NOTICE file. + + 7. Disclaimer of Warranty. Unless required by applicable law or + agreed to in writing, Licensor provides the Work (and each + Contributor provides its Contributions) on an "AS IS" BASIS, + WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or + implied, including, without limitation, any warranties or conditions + of TITLE, NON-INFRINGEMENT, MERCHANTABILITY, or FITNESS FOR A + PARTICULAR PURPOSE. You are solely responsible for determining the + appropriateness of using or redistributing the Work and assume any + risks associated with Your exercise of permissions under this License. + + 8. Limitation of Liability. In no event and under no legal theory, + whether in tort (including negligence), contract, or otherwise, + unless required by applicable law (such as deliberate and grossly + negligent acts) or agreed to in writing, shall any Contributor be + liable to You for damages, including any direct, indirect, special, + incidental, or consequential damages of any character arising as a + result of this License or out of the use or inability to use the + Work (including but not limited to damages for loss of goodwill, + work stoppage, computer failure or malfunction, or any and all + other commercial damages or losses), even if such Contributor + has been advised of the possibility of such damages. + + 9. Accepting Warranty or Additional Liability. While redistributing + the Work or Derivative Works thereof, You may choose to offer, + and charge a fee for, acceptance of support, warranty, indemnity, + or other liability obligations and/or rights consistent with this + License. However, in accepting such obligations, You may act only + on Your own behalf and on Your sole responsibility, not on behalf + of any other Contributor, and only if You agree to indemnify, + defend, and hold each Contributor harmless for any liability + incurred by, or claims asserted against, such Contributor by reason + of your accepting any such warranty or additional liability. + + END OF TERMS AND CONDITIONS + + APPENDIX: How to apply the Apache License to your work. + + To apply the Apache License to your work, attach the following + boilerplate notice, with the fields enclosed by brackets "[]" + replaced with your own identifying information. (Don't include + the brackets!) The text should be enclosed in the appropriate + comment syntax for the file format. We also recommend that a + file or class name and description of purpose be included on the + same "printed page" as the copyright notice for easier + identification within third-party archives. + + Copyright Zhengxiao Du + + 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. \ No newline at end of file diff --git a/PyTorch/built-in/foundation/CodeGeeX2/MODEL_LICENSE b/PyTorch/built-in/foundation/CodeGeeX2/MODEL_LICENSE new file mode 100644 index 0000000000000000000000000000000000000000..a33932b44d906e30f77684d510f6eb3cabb4e5f9 --- /dev/null +++ b/PyTorch/built-in/foundation/CodeGeeX2/MODEL_LICENSE @@ -0,0 +1,33 @@ +The CodeGeeX License + +1. Definitions + +“Licensor” means the CodeGeeX Model Team that distributes its Software. + +“Software” means the CodeGeeX model parameters made available under this license. + +2. License Grant + +Subject to the terms and conditions of this License, the Licensor hereby grants to you a non-exclusive, worldwide, non-transferable, non-sublicensable, revocable, royalty-free copyright license to use the Software solely for your non-commercial research purposes. + +The above copyright notice and this permission notice shall be included in all copies or substantial portions of the Software. + +3. Restriction + +You will not use, copy, modify, merge, publish, distribute, reproduce, or create derivative works of the Software, in whole or in part, for any commercial, military, or illegal purposes. + +You will not use the Software for any act that may undermine China's national security and national unity, harm the public interest of society, or infringe upon the rights and interests of human beings. + +4. Disclaimer + +THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY, FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM, OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE SOFTWARE. + +5. Limitation of Liability + +EXCEPT TO THE EXTENT PROHIBITED BY APPLICABLE LAW, IN NO EVENT AND UNDER NO LEGAL THEORY, WHETHER BASED IN TORT, NEGLIGENCE, CONTRACT, LIABILITY, OR OTHERWISE WILL ANY LICENSOR BE LIABLE TO YOU FOR ANY DIRECT, INDIRECT, SPECIAL, INCIDENTAL, EXEMPLARY, OR CONSEQUENTIAL DAMAGES, OR ANY OTHER COMMERCIAL LOSSES, EVEN IF THE LICENSOR HAS BEEN ADVISED OF THE POSSIBILITY OF SUCH DAMAGES. + +6. Dispute Resolution + +This license shall be governed and construed in accordance with the laws of People’s Republic of China. Any dispute arising from or in connection with this License shall be submitted to Haidian District People's Court in Beijing. + +Note that the license is subject to update to a more comprehensive version. For any questions related to the license and copyright, please contact us at license@zhipuai.cn. \ No newline at end of file diff --git a/PyTorch/built-in/foundation/CodeGeeX2/README.md b/PyTorch/built-in/foundation/CodeGeeX2/README.md new file mode 100644 index 0000000000000000000000000000000000000000..cdf7bbb9af665732b1f93252be3b59dd6c05a99b --- /dev/null +++ b/PyTorch/built-in/foundation/CodeGeeX2/README.md @@ -0,0 +1,334 @@ +# codegeex2_NPU + +- [概述](概述.md) +- [准备训练环境](准备训练环境.md) +- [开始训练](开始训练.md) +- [训练结果展示](训练结果展示.md) +- [版本说明](版本说明.md) + + + +# 概述 + +## 简述 + +CodeGeeX2 是多语言代码生成模型 [CodeGeeX](https://github.com/THUDM/CodeGeeX) ([KDD’23](https://arxiv.org/abs/2303.17568)) 的第二代模型。不同于一代 CodeGeeX（完全在国产华为昇腾芯片平台训练） ，CodeGeeX2 是基于 [ChatGLM2](https://github.com/THUDM/ChatGLM2-6B) 架构加入代码预训练实现，得益于 ChatGLM2 的更优性能，CodeGeeX2 在多项指标上取得性能提升（+107% > CodeGeeX；仅60亿参数即超过150亿参数的 StarCoder-15B 近10%），更多特性包括： + +- **更强大的代码能力**：基于 ChatGLM2-6B 基座语言模型，CodeGeeX2-6B 进一步经过了 600B 代码数据预训练，相比一代模型，在代码能力上全面提升，[HumanEval-X](https://huggingface.co/datasets/THUDM/humaneval-x) 评测集的六种编程语言均大幅提升 (Python +57%, C++ +71%, Java +54%, JavaScript +83%, Go +56%, Rust +321%)，在Python上达到 35.9% 的 Pass@1 一次通过率，超越规模更大的 StarCoder-15B。 +- **更优秀的模型特性**：继承 ChatGLM2-6B 模型特性，CodeGeeX2-6B 更好支持中英文输入，支持最大 8192 序列长度，推理速度较一代 CodeGeeX-13B 大幅提升，量化后仅需6GB显存即可运行，支持轻量级本地化部署。 +- **更全面的AI编程助手**：CodeGeeX插件（[VS Code](https://marketplace.visualstudio.com/items?itemName=aminer.codegeex), [Jetbrains](https://plugins.jetbrains.com/plugin/20587-codegeex)）后端升级，支持超过100种编程语言，新增上下文补全、跨文件补全等实用功能。结合 Ask CodeGeeX 交互式AI编程助手，支持中英文对话解决各种编程问题，包括且不限于代码解释、代码翻译、代码纠错、文档生成等，帮助程序员更高效开发。 +- **更开放的协议**：CodeGeeX2-6B 权重对学术研究完全开放，填写[登记表](https://open.bigmodel.cn/mla/form?mcode=CodeGeeX2-6B)申请商业使用。 + + + +- 参考实现： + + ``` + url=https://github.com/THUDM/ChatGLM2-6B + commit_id=877ef10d85c93ddfcfe945fcdc764393a52541b8 + url=https://huggingface.co/THUDM/chatglm2-6b/tree/v1.0 + commit_id=0ade0d38ac00258ae09450696315c2ff0b1faf12 + ``` + +- 适配昇腾 AI 处理器的实现： + + ``` + url=https://gitee.com/ascend/ModelZoo-PyTorch.git + code_path=PyTorch/built-in/foundation + ``` + + +# 准备训练环境 + +## 准备环境 + +默认配置需要每张卡有60G以上空闲内存。 + +- 当前模型支持的 PyTorch 版本和已知三方库依赖如下表所示。 + + **表 1** 版本支持表 + + | Torch_Version | 三方库依赖版本 | + | :-----------: | :-------------: | + | PyTorch 1.11 | deepspeed 0.9.2 | + +- 环境准备指导。 + + 请参考《[Pytorch框架训练环境准备](https://www.hiascend.com/document/detail/zh/ModelZoo/pytorchframework/ptes)》。 + +- 安装依赖。 + + 1. 安装基础依赖 + + 在模型源码包根目录下执行命令，安装模型对应PyTorch版本需要的依赖。 + + ``` + pip install -r requirements.txt # PyTorch1.11版本 + ``` + + 2. 安装deepspeed_npu插件 + + ``` + # v0.9.2分支 + pip3 install deepspeed==0.9.2 + git clone https://gitee.com/ascend/DeepSpeed.git -b v0.9.2 deepspeed_npu + cd deepspeed_npu + git checkout 5c7c89930f0b70ea586d5db63f8e66477d5d9d9f + pip3 install . + ``` + + 3. 替换transformers依赖文件 + + ``` + # 自动替换无法替换三方库中的文件。 + pip show transformers + # 获取transformers的Location路径 + # 使用fix文件夹下的tranining_args.py替换路径下transformers/tranining_args.py + 规避保存权重问题： + # 使用fix文件夹下的modeling_utils.py替换transformers/modeling_utils.py + ``` + + +## 准备数据集 + +1. 获取数据集。 + + code_alpaca_20k + + + + ``` + ├── ptuning + ├──train.json + ├──dev.json + ``` + +2. 预处理数据集。 + 为了方便启动训练后，不用再每次重复加载处理数据集，故提前进行处理。也可以下载提前处理好的[数据集](待补充） + +```shell +bash preprocess.sh +``` + +处理好的数据集位于同目录下的train_datasets文件夹下，参考目录如下 + +``` + ├── train_datasets + ├──data-00000-of-00003.arrow + ├──data-00001-of-00003.arrow + ├──data-00002-of-00003.arrow + ├──dataset_info.json + ├──state.json +``` + + +## 准备模型权重 + +1. 获取语言识别模型和预训练权重 + + 1)用户从[链接]([THUDM/codegeex2-6b at main (huggingface.co)](https://huggingface.co/THUDM/codegeex2-6b/tree/main))自行获取模型文件（除了modeling_chatglm.py）和权重文件（pytorch_model-0000*-of-00007.bin），并放于model目录下，微调依赖该模型权重。 + + 2)Please Do NOT overwrite modeling_chatglm.py + + The "model" directory is as follows + + ```shell + ├── model + ├──config.json + ├──configuration_chatglm.py + ├──ice_text.model + ├──pytorch_model-00001-of-00007.bin + ├──pytorch_model-00002-of-00007.bin + ├──pytorch_model-00003-of-00007.bin + ├──pytorch_model-00004-of-00007.bin + ├──pytorch_model-00005-of-00007.bin + ├──pytorch_model-00006-of-00007.bin + ├──pytorch_model-00007-of-00007.bin + ├──pytorch_model.bin.index.json + ├──quantization.py + ├──test_modeling_chatglm.py + ├──tokenization_chatglm.py + ├──tokenizer_config.json + ├──tokenizer.model + ├──modeling_chatglm.py + ``` + + + +# 开始训练 + +## 训练模型 + +1. 进入解压后的源码包根目录。 + + ``` + cd /${模型文件夹名称}/ptuning + ``` + + 修改ptuning目录下的env_npu.sh，修改引用的环境变量位置。 + +2. 运行训练脚本。 + + 该模型P-Tuning v2支持单机单卡，全参数fintune支持单机8卡。 + + - P-Tuning v2 + + 启动P-Tuning v2。 + + ``` + bash train.sh + ``` + + - 全参数finetune + + 启动8卡微调。 + 可以用deepspeed.json配置deepspeed参数，目前默认使用zero2 + + ``` + bash ds_train_fintune.sh + ``` + + + 模型训练参数说明如下。 + + ``` + 公共参数： + --max_source_length //处理后句子长度 + --max_target_length //目标数据长度 + --per_device_train_batch_size //每卡训练批次大小 + --gradient_accumulation_steps //梯度更新步数 + --max_steps //最大训练步数 + --logging_steps //打印信息步数 + --save_steps //保存参数步数 + --learning_rate //学习率 + ``` + + 训练完成后，权重文件保存在当前路径下，并输出模型训练相关信息。 + +## 验证模型 + +1. 全参数finetune验证 + + 运行以下命令 + + ``` + cd /${模型文件夹名称}/scripts + ./run_humanevalx.sh + ``` + + 生成结果在屏幕上显示 + +# 训练结果展示 + +**表 1** 训练结果展示表 + + +| NAME | SamplesPerSec | Iterations | DataType | Torch_Version | Card | +| :-----------: | :-----------: | :--------: | :------: | :-----------: | :--: | +| Finetune -NPU | (待补充) | (待补充) | bf16? | 1.11 | 910B | +| Finetune -GPU | (待补充) | (待补充) | bf16? | 1.11 | A800 | + +说明：P-Tuning 仅打通功能，无性能优化。 + +**表 2** 评估结果展示表 + +| 评估项 | NPU | GPU | +| :-----: | :--: | :--: | +| human pass@1 | 0.37 | 0.35 | + + +说明：该结果是step=100的验证结果。 + +# 版本说明 + +## 变更 + +2023.9.1：首次发布。 + +## FAQ + +1. 报错提示deepspeed.py需要版本大于等于0.6.5 + + ``` + # 关闭版本检测（如安装0.9.2版本无需此操作） + # 若遇到该报错 + pip show transformers + # 复制Location路径 + # 使用fix文件夹下的deepspeed.py替换路径下transformers/deepspeed.py + ``` + +2. 加载参数阶段有卡死现象 + + ``` + 删除root下的cache目录，重新运行 + ``` + +3. 单卡阶段报embedding_dense_grad算子错误 + + ``` + enbedding当前版本，不支持动静合一，静态有部分shape不支持,新版本已修复 + # 若遇到该报错 + 修改main.py文件 + torch.npu.set_compile_mode(jit_compile=False) + ``` + +4. 提示so文件错误 + + ``` + 提示so文件找不到 + # 若遇到该报错 + 全局搜索so的位置，然后导入环境变量 + export LD_LIBRARY_PATH=/usr/:$LD_LIBRARY_PATH + ``` + +5. eval提示scaledsoftmax报错 + + ``` + 算子shape泛化性还有问题 + # 若遇到该报错 + 搜索output文件夹生成的modeling_chatglm.py文件， + self.scale_mask_softmax 设置为false + ``` + +6. 遇到ImportError:/root/miniconda3/envs/codegeex/bin/../lib/libgomp.so.1: cannot allocate memory in static TLS block + + ``` + export LD_PRELOAD=/root/miniconda3/envs/codegeex/bin/../lib/libgomp.so.1 + ``` + +7. 微调时出现AttributeError或RuntimeError + +module 'torch_npu' has no attribute 'npu_rotary_mul' 或 + +RuntimeError:Error!, The last dimension of input tensor shoule be within the range of [32,2048] and be divisible by32 + +``` +修改modeling_chatglm.py文件: +USE_NPU_ROTARY=False +USE_SCALED_SOFTMAX=False +``` + +PS: 设置为True能提高性能 + +8. 如果cann不支持flash_attention + +报错提示为module 'torch_npu' has no attribute 'npu_flash_attention' + +``` +修改modeling_chatglm.py文件: +USE_FLASH=False +``` + +​ PS: 设置为True能提高性能 + +9. 如果cann不支持npu_rms_norm + +报错提示为module 'torch_npu' has no attribute 'npu_rms_norm' + +``` +修改modeling_chatglm.py文件: +USE_NPU_RMSNORM=False +``` + +​ PS: 设置为True能提高性能 + diff --git a/PyTorch/built-in/foundation/CodeGeeX2/api.py b/PyTorch/built-in/foundation/CodeGeeX2/api.py new file mode 100644 index 0000000000000000000000000000000000000000..e4236b74433af8f337a4224bae2e81d280156212 --- /dev/null +++ b/PyTorch/built-in/foundation/CodeGeeX2/api.py @@ -0,0 +1,60 @@ +from fastapi import FastAPI, Request +from transformers import AutoTokenizer, AutoModel +import uvicorn, json, datetime +import torch + +DEVICE = "cuda" +DEVICE_ID = "0" +CUDA_DEVICE = f"{DEVICE}:{DEVICE_ID}" if DEVICE_ID else DEVICE + + +def torch_gc(): + if torch.cuda.is_available(): + with torch.cuda.device(CUDA_DEVICE): + torch.cuda.empty_cache() + torch.cuda.ipc_collect() + + +app = FastAPI() + + +@app.post("/") +async def create_item(request: Request): + global model, tokenizer + json_post_raw = await request.json() + json_post = json.dumps(json_post_raw) + json_post_list = json.loads(json_post) + prompt = json_post_list.get('prompt') + history = json_post_list.get('history') + max_length = json_post_list.get('max_length') + top_p = json_post_list.get('top_p') + temperature = json_post_list.get('temperature') + response, history = model.chat(tokenizer, + prompt, + history=history, + max_length=max_length if max_length else 2048, + top_p=top_p if top_p else 0.7, + temperature=temperature if temperature else 0.95) + now = datetime.datetime.now() + time = now.strftime("%Y-%m-%d %H:%M:%S") + answer = { + "response": response, + "history": history, + "status": 200, + "time": time + } + log = "[" + time + "] " + '", prompt:"' + prompt + '", response:"' + repr(response) + '"' + print(log) + torch_gc() + return answer + + +if __name__ == '__main__': + tokenizer = AutoTokenizer.from_pretrained("THUDM/chatglm2-6b", trust_remote_code=True) + model = AutoModel.from_pretrained("THUDM/chatglm2-6b", trust_remote_code=True).cuda() + # 多显卡支持，使用下面三行代替上面两行，将num_gpus改为你实际的显卡数量 + # model_path = "THUDM/chatglm2-6b" + # tokenizer = AutoTokenizer.from_pretrained(model_path, trust_remote_code=True) + # model = load_model_on_gpus(model_path, num_gpus=2) + model.eval() + uvicorn.run(app, host='0.0.0.0', port=8000, workers=1) diff --git a/PyTorch/built-in/foundation/CodeGeeX2/benchmark/README.md b/PyTorch/built-in/foundation/CodeGeeX2/benchmark/README.md new file mode 100644 index 0000000000000000000000000000000000000000..ec0391c05b7a4aaf37219577d479d41be8cd5d1c --- /dev/null +++ b/PyTorch/built-in/foundation/CodeGeeX2/benchmark/README.md @@ -0,0 +1,75 @@ +# HumanEval-X: A new benchmark for Multilingual Program Synthesis + +🌐 中文 + +HumanEval-X is a new benchmark for better evaluating the multilingual ability of code generation models. While previous works evaluate multilingual program synthesis under semantic similarity (e.g., [CodeBLEU](https://arxiv.org/abs/2009.10297)) which is often misleading, HumanEval-X evaluates the functional correctness of the generated programs. HumanEval-X consists of 820 high-quality human-crafted data samples (each with test cases) in Python, C++, Java, JavaScript, and Go, and can be used for various tasks. + + + +

An illustration of tasks supported by HumanEval-X. Declarations, docstrings, and solutions are marked with red, green, and blue respectively. Code generation uses declaration and docstring as input, to generate solution. Code translation uses declaration in both languages and translate the solution in source language to the one in target language.

+ +In HumanEval-X, every sample in each language contains declaration, docstring, and solution, which can be combined in various ways to support different downstream tasks including generation, translation, summarization, etc. We currently focus on two tasks: **code generation** and **code translation**. For code generation, the model uses declaration and docstring as input to generate the solution. For code translation, the model uses declarations in both languages and the solution in the source language as input, to generate solutions in the target language. We remove the description during code translation to prevent the model from directly solving the problem. For both tasks, we use the unbiased pass@k metric proposed in [Codex](https://arxiv.org/abs/2107.03374): $\text{pass}@k:= \mathbb{E}[1-\frac{\tbinom{n-c}{k}}{\tbinom{n}{k}}]$, with $n=200$ and $k\in(1,10,100)$. + +## How to use HumanEval-X + +Data are stored in ``codegeex/benchmark/humaneval-x/[LANG]/data/humaneval_[LANG].jsonl.gz``, using JSON list format. There are six keys: + +* ``task_id``: indicates the target language and ID of the problem. Language is one of ["Python", "Java", "JavaScript", "CPP", "Go"]. +* ``prompt``: the function declaration and docstring, used for code generation. +* ``declaration``: only the function declaration, used for code translation. +* ``canonical_solution``: human-crafted example solutions. +* ``test``: hidden test samples, used for evaluation. +* ``example_test``: public test samples (appeared in prompt), used for evaluation. + +### Evaluation Environment + +The evaluation of the generated codes involves compiling and running in multiple programming languages. The versions of the programming language environments and packages we use are as follows: + +| Dependency | Version | +| ---------- | -------- | +| Python | 3.8.12 | +| JDK | 18.0.2.1 | +| Node.js | 16.14.0 | +| js-md5 | 0.7.3 | +| C++ | 11 | +| g++ | 7.5.0 | +| Boost | 1.71.0 | +| OpenSSL | 3.0.0 | +| go | 1.18.4 | + +In order to save everyone the trouble of setting up the environments for these languages, we build a Docker image with the required environments and CodeGeeX installed. + +You can directly pull the image from Docker Hub: + +```bash +docker pull rishubi/codegeex:latest +``` + +Alternatively, if you are familiar with Dockerfile, you can build the image from `codegeex/docker/Dockerfile` or configure the Dockerfile as you like it: + +```bash +cd codegeex/docker +docker build [OPTIONS] . +``` + +After obtaining the image, you can build a container using the following command: + +```bash +docker run -it --gpus all --mount type=bind,source=,target= [OPTIONS] +``` + +### Evaluation + +We recommend evaluating in [the provided image](#evaluation-environment). To evaluate the generated samples, save generated codes in the following JSON list format: + +``` +{"task_id": "../..", "generation: "..."} +{"task_id": "../..", "generation: "..."} +... +``` + +and evaluate them using the following script under the root directory of the repository (please execute with caution, the generated codes might have unexpected behaviours though with very low possibility. See the warnings in [execution.py](execution.py) and uncomment the execution lines at your own risk): + +```bash +bash scripts/evaluate_humaneval_x.sh +``` diff --git a/PyTorch/built-in/foundation/CodeGeeX2/benchmark/README_zh.md b/PyTorch/built-in/foundation/CodeGeeX2/benchmark/README_zh.md new file mode 100644 index 0000000000000000000000000000000000000000..ef07392120d3dc76c7a1c6d3347a1a9ae6fb7db4 --- /dev/null +++ b/PyTorch/built-in/foundation/CodeGeeX2/benchmark/README_zh.md @@ -0,0 +1,75 @@ +# HumanEval-X: 多语言代码生成基准 + +🌐 English + +为了更好地评测代码生成模型的多语言生成能力，我们构建了一个新基准HumanEval-X。此前，多语言代码生成能力是基于语义相似度（比如[CodeBLEU](https://arxiv.org/abs/2009.10297)）衡量的，具有一定误导性；HumanEval-X则可用于衡量生成代码的功能正确性。HumanEval-X包含820个高质量手写样本，覆盖Python、C++、Java、JavaScript、Go，可用于多种任务。 + + + +

HumanEval-X支持的任务示例。声明、描述、解答分别用红、绿、蓝色标注。代码生成将声明与描述作为输入，输出解答。代码翻译将两种语言的声明与源语言的解答作为输入，输出目标语言的解答。

+ +HumanEval-X中每个语言的样本，包含了声明、描述和解答，它们之间的组合可以支持不同的下游任务，包括生成、翻译、概括等。我们目前关注两个任务：**代码生成**与**代码翻译**。对于代码生成任务，模型将函数声明与文档字符串作为输入，输出函数实现；对于代码翻译任务，模型将两种语言的函数声明与源语言的实现作为输入，输出目标语言上的实现。我们在代码翻译任务中不将文档字符串输入模型，以避免模型直接通过描述生成答案。在两种任务下，我们都采用[Codex](https://arxiv.org/abs/2107.03374)所使用的无偏pass@k指标：$\text{pass}@k:= \mathbb{E}[1-\frac{\tbinom{n-c}{k}}{\tbinom{n}{k}}]$, $n=200$, $k\in(1,10,100)$。 + +## 如何使用HumanEval-X + +样本使用JSON列表格式存储在``codegeex/benchmark/humaneval-x/[LANG]/data/humaneval_[LANG].jsonl.gz``，每条样本包含6个部分： + +* ``task_id``: 题目的目标语言与ID。语言为["Python", "Java", "JavaScript", "CPP", "Go"]中之一。 +* ``prompt``: 函数声明与描述，用于代码生成。 +* ``declaration``: 仅有函数声明，用于代码翻译。 +* ``canonical_solution``: 手写的示例解答。 +* ``test``: 隐藏测例，用于评测。 +* ``example_test``: 提示中出现的公开测例，用于评测。 + +### 评测环境 + +评测生成的代码需要使用多种语言编译、运行。我们使用的各编程语言依赖及所用包的版本如下： + +| 依赖 | 版本 | +| ------- | -------- | +| Python | 3.8.12 | +| JDK | 18.0.2.1 | +| Node.js | 16.14.0 | +| js-md5 | 0.7.3 | +| C++ | 11 | +| g++ | 7.5.0 | +| Boost | 1.71.0 | +| OpenSSL | 3.0.0 | +| go | 1.18.4 | + +为了省去使用者配置这些语言环境的麻烦，我们构建了一个Docker镜像，并在其中配置了所需要的环境。 + +可以直接从Docker Hub拉取镜像： + +```bash +docker pull rishubi/codegeex:latest +``` + +如果您熟悉Dockerfile，也可以从`codegeex/docker/Dockerfile`构建镜像，或者修改之以定制自己的配置： + +```bash +cd codegeex/docker +docker build [OPTIONS] . +``` + +获取镜像后，使用如下命令创建容器： + +```bash +docker run -it --gpus all --mount type=bind,source=,target= [OPTIONS] +``` + +### 评测 + +我们推荐使用给定的[评测环境](#评测环境)进行评测。在评测前，将生成的代码以如下JSON列表形式存储： + +``` +{"task_id": "../..", "generation: "..."} +{"task_id": "../..", "generation: "..."} +... +``` + +并在本仓库的根目录下使用如下指令（请谨慎执行，生成的代码可能有极低概率产生意外行为。在[execution.py](execution.py)中查看警告并取消执行代码的注释，风险自负）： + +```bash +bash scripts/evaluate_humaneval_x.sh +``` diff --git a/PyTorch/built-in/foundation/CodeGeeX2/benchmark/__init__.py b/PyTorch/built-in/foundation/CodeGeeX2/benchmark/__init__.py new file mode 100644 index 0000000000000000000000000000000000000000..e69de29bb2d1d6434b8b29ae775ad8c2e48c5391 diff --git a/PyTorch/built-in/foundation/CodeGeeX2/benchmark/evaluate_humaneval_x.py b/PyTorch/built-in/foundation/CodeGeeX2/benchmark/evaluate_humaneval_x.py new file mode 100644 index 0000000000000000000000000000000000000000..076e0e89377192bb065bcff12e3f603b36f8111f --- /dev/null +++ b/PyTorch/built-in/foundation/CodeGeeX2/benchmark/evaluate_humaneval_x.py @@ -0,0 +1,239 @@ +import os +import sys +import fire +import json +import gzip +import regex +import numpy as np + +from typing import * +from tqdm.auto import tqdm +from collections import defaultdict +from concurrent.futures import ThreadPoolExecutor, as_completed + +from codegeex.benchmark.utils import read_dataset, IMPORT_HELPER +from codegeex.benchmark.metric import estimate_pass_at_k +from codegeex.benchmark.execution import check_correctness + +LANGUAGE_NAME = { + "cpp" : "CPP", + "go" : "Go", + "java" : "Java", + "js" : "JavaScript", + "python": "Python", +} + + +def process_humaneval_test(sample, problems, example_test=False): + task_id = sample["task_id"] + language = task_id.split("/")[0].lower() + + prompt = sample["prompt"] + if example_test and "example_test" in problems[task_id] and problems[task_id]["example_test"] != "": + test = problems[task_id]["example_test"] + else: + test = problems[task_id]["test"] + code = sample["generation"] + + # Pre-process for different languages + if language == "python": + code_ = [] + for line in code.split("\n"): + if (len(line.strip()) > 0 and line[0] != ' ' and line[0] != '\t'): + break + code_.append(line) + code = "\n".join(code_) + test_setup = "\n".join(IMPORT_HELPER["python"]) + "\n" + test_string = test_setup + prompt + code + "\n" + test + "\n" + elif language == "cpp": + test_set_up = "" + for s in IMPORT_HELPER["cpp"]: + if s not in prompt: + test_set_up += s + "\n" + test_string = test_set_up + "\n" + prompt + code + "\n" + test + elif language == "java": + test_string = prompt + code + "\n" + test + elif language == "js" or language == "javascript": + test_string = prompt + code + "\n" + test + elif language == "go": + import_string = problems[task_id]["import"] + prompt = prompt.replace(import_string, "") + if example_test and "example_test" in problems[task_id]: + test = problems[task_id]["example_test"] + else: + test = problems[task_id]["test"] + test_setup = problems[task_id]["test_setup"] + other_pkgs = [] + for pkg in IMPORT_HELPER["go"]: + if pkg not in test_setup: + p = pkg.split("/")[-1] + if p + "." in code: + other_pkgs.append(f"\"{pkg}\"") + if other_pkgs: + import_other_pkgs = "import (\n" + " ".join([p + "\n" for p in other_pkgs]) + ")" + test_string = test_setup + "\n" + import_other_pkgs + "\n" + prompt + code + "\n" + test + else: + test_string = test_setup + "\n" + prompt + code + "\n" + test + elif language == "rust": + main = "\nfn main(){ \n } \n" + declaration = problems[task_id]["declaration"] + test_string = main + declaration + prompt + code + test + + return test_string + + +def stream_jsonl_all(filename: str) -> Iterable[Dict]: + results = [] + if filename.endswith(".gz"): + fp = gzip.open(open(filename, "rb"), "rt") + else: + fp = open(filename, "r") + for line in fp: + if any(not x.isspace() for x in line): + results.append(json.loads(line)) + fp.close() + + return results + + +def evaluate_functional_correctness( + input_file: str = None, + tmp_dir: str = "./", + n_workers: int = 32, + timeout: float = 500.0, + problem_file: str = "../data/humaneval_python.jsonl.gz", + out_dir: str = None, + k: List[int] = [1, 10, 100], + test_groundtruth: bool = False, + example_test: bool = False, +): + if example_test: + print("Example test...") + + problems = read_dataset(problem_file, + dataset_type="humaneval") + sample_jsonl = stream_jsonl_all(input_file) + + if example_test: + suffix = "_example_test.jsonl" + else: + suffix = "_results.jsonl" + if out_dir is not None: + if not os.path.exists(out_dir): + os.makedirs(out_dir) + out_file = os.path.join(out_dir, input_file.split('/')[-1].replace(".jsonl", suffix)) + else: + out_file = os.path.join(input_file.replace(".jsonl", suffix)) + + if "/codegeex/benchmark/humaneval-x/" in input_file: + test_groundtruth = True + + if "-to-" in input_file: + translation_mode = True + else: + translation_mode = False + + with ThreadPoolExecutor(max_workers=n_workers) as executor: + + futures = [] + completion_id = Counter() + n_samples = 0 + results = defaultdict(list) + + if test_groundtruth: + print("Testing ground truth...") + for sample in tqdm(problems.values()): + task_id = sample["task_id"] + lang = task_id.split("/")[0].lower() + if lang == "javascript": + lang = "js" + tmp_dir_ = os.path.join(tmp_dir, lang, "evaluation") + sample["generation"] = sample["canonical_solution"] + sample["test_code"] = process_humaneval_test(sample, problems, example_test) + if sample["test_code"] is None: + continue + args = (task_id, sample, lang, timeout, tmp_dir_, completion_id[task_id]) + future = executor.submit(check_correctness, *args) + futures.append(future) + completion_id[task_id] += 1 + n_samples += 1 + else: + print("Reading samples...") + for sample in tqdm(sample_jsonl): + task_id = sample["task_id"] + lang = task_id.split("/")[0].lower() + if translation_mode: + task_id = sample["task_id"].split("/")[-1] + lang = regex.findall("-to-.*-", input_file)[0].split("-to-")[-1].rstrip("-") + for l in LANGUAGE_NAME: + if l in lang: + lang = l + break + task_id = f"{LANGUAGE_NAME[lang]}/{task_id}" + if lang == "javascript": + lang = "js" + tmp_dir_ = os.path.join(tmp_dir, lang, "evaluation") + sample["task_id"] = task_id + sample["test_code"] = process_humaneval_test(sample, problems, example_test) + if sample["test_code"] is None: + continue + if "completion_id" in sample: + completion_id_ = sample["completion_id"] + else: + completion_id_ = completion_id[task_id] + args = (task_id, sample, lang, timeout, tmp_dir_, completion_id_) + future = executor.submit(check_correctness, *args) + futures.append(future) + completion_id[task_id] += 1 + n_samples += 1 + + print(completion_id) + if len(completion_id) == len(problems): + evaluate_pass_at_k = True + else: + evaluate_pass_at_k = False + + print("Running test suites...") + for future in tqdm(as_completed(futures), total=len(futures)): + result = future.result() + results[result["task_id"]].append((result["completion_id"], result)) + + # Calculate pass@k. + total, correct = [], [] + for result in results.values(): + passed = [r[1]["passed"] for r in result] + total.append(len(passed)) + correct.append(sum(passed)) + total = np.array(total) + correct = np.array(correct) + if evaluate_pass_at_k: + ks = k + pass_at_k = {f"pass@{k}": estimate_pass_at_k(total, correct, k).mean() + for k in ks if (total >= k).all()} + print(pass_at_k) + else: + print("Total:", np.sum(total)) + print("Correct:", np.sum(correct)) + + print("Writing to: ", out_file) + if out_file.endswith(".gz"): + fp = gzip.GzipFile(fileobj=open(out_file, "wb"), mode="wb") + for res in results.values(): + for r in res: + fp.write((json.dumps(r[1]) + "\n").encode("utf-8")) + else: + fp = open(out_file, 'w') + for res in results.values(): + for r in res: + fp.write(json.dumps(r[1]) + "\n") + fp.close() + + print("Evaluation finished.") + + +def main(): + fire.Fire(evaluate_functional_correctness) + + +if __name__ == "__main__": + sys.exit(main()) diff --git a/PyTorch/built-in/foundation/CodeGeeX2/benchmark/execution.py b/PyTorch/built-in/foundation/CodeGeeX2/benchmark/execution.py new file mode 100644 index 0000000000000000000000000000000000000000..cbdf14f4954d9f458d1e2c8e5fb8016865d6244a --- /dev/null +++ b/PyTorch/built-in/foundation/CodeGeeX2/benchmark/execution.py @@ -0,0 +1,549 @@ +import contextlib +import faulthandler +import io +import multiprocessing +import os +import platform +import signal +import random +import subprocess +import tempfile +import gzip +import json +from typing import * + +def dicts_to_jsonl(data_list: list, filename: str, compress: bool = True) -> None: + """ + Method saves list of dicts into jsonl file. + :param data: (list) list of dicts to be stored, + :param filename: (str) path to the output file. If suffix .jsonl is not given then methods appends + .jsonl suffix into the file. + :param compress: (bool) should file be compressed into a gzip archive? + """ + sjsonl = '.jsonl' + sgz = '.gz' + # Check filename + if not filename.endswith(sjsonl): + filename = filename + sjsonl + # Save data + + if compress: + filename = filename + sgz + with gzip.open(filename, 'w') as compressed: + for ddict in data_list: + jout = json.dumps(ddict) + '\n' + jout = jout.encode('utf-8') + compressed.write(jout) + else: + with open(filename, 'w') as out: + for ddict in data_list: + jout = json.dumps(ddict) + '\n' + out.write(jout) + + +def check_correctness( + task_id: str, + sample: dict, + language_type: str, + timeout: float = 3.0, + tmp_dir: str = None, + completion_id: Optional[int] = None, +) -> Dict: + """ + Evaluates the functional correctness of a completion by running the test + suite provided in the problem. + """ + + def unsafe_execute(tmp_dir): + random_id = random.uniform(1, 1000) + if "python" in language_type.lower(): + with create_tempdir(): + + # These system calls are needed when cleaning up tempdir. + import os + import shutil + rmtree = shutil.rmtree + rmdir = os.rmdir + chdir = os.chdir + + # Disable functionalities that can make destructive changes to the test. + reliability_guard() + + try: + exec_globals = {} + with swallow_io(): + with time_limit(timeout): + # WARNING + # This program exists to execute untrusted model-generated code. Although + # it is highly unlikely that model-generated code will do something overtly + # malicious in response to this test suite, model-generated code may act + # destructively due to a lack of model capability or alignment. + # Users are strongly encouraged to sandbox this evaluation suite so that it + # does not perform destructive actions on their host or network. + # Once you have read this disclaimer and taken appropriate precautions, + # uncomment the following line and proceed at your own risk: + exec(sample["test_code"], exec_globals) + result.append("passed") + except TimeoutException: + result.append("timed out") + except AssertionError as e: + result.append(f"failed: AssertionError") + except BaseException as e: + result.append(f"failed: {e}") + + # Needed for cleaning up. + shutil.rmtree = rmtree + os.rmdir = rmdir + os.chdir = chdir + + elif "go" in language_type.lower(): + assert tmp_dir is not None, "Go should be evaluated in a dir where necessary module files installed." + + import os + import shutil + + if "tmp" not in tmp_dir: + tmp_dir = os.path.join(tmp_dir, "tmp") + tmp_dir = os.path.join(tmp_dir, f"{task_id.replace('/', '-')}-{random_id}") + if not os.path.exists(tmp_dir): + os.makedirs(tmp_dir) + + os.chdir(tmp_dir) + open(f"main_test.go", 'w').write(sample["test_code"]) + try: + exec_result = None + with time_limit(timeout): + # WARNING + # This program exists to execute untrusted model-generated code. Although + # it is highly unlikely that model-generated code will do something overtly + # malicious in response to this test suite, model-generated code may act + # destructively due to a lack of model capability or alignment. + # Users are strongly encouraged to sandbox this evaluation suite so that it + # does not perform destructive actions on their host or network. + # Once you have read this disclaimer and taken appropriate precautions, + # uncomment the following line and proceed at your own risk: + exec_result = subprocess.run(["go", "test", f"-timeout={timeout}s", "main_test.go"], timeout=timeout, capture_output=True) + + if exec_result.returncode == 0: + result.append("passed") + else: + if exec_result.stderr: + try: + err = exec_result.stderr.decode() + except: + err = exec_result.stderr + else: + try: + err = exec_result.stdout.decode() + except: + err = exec_result.stdout + result.append(f"failed: {err}") + + except TimeoutException: + result.append("timed out") + + shutil.rmtree(tmp_dir) + elif "js" in language_type.lower(): + import os + import shutil + + if "tmp" not in tmp_dir: + tmp_dir = os.path.join(tmp_dir, "tmp") + tmp_dir = os.path.join(tmp_dir, f"{task_id.replace('/', '-')}-{random_id}") + if not os.path.exists(tmp_dir): + os.makedirs(tmp_dir) + + os.chdir(tmp_dir) + open(f"test.js", 'w').write(sample["test_code"]) + try: + exec_result = None + with time_limit(timeout): + # WARNING + # This program exists to execute untrusted model-generated code. Although + # it is highly unlikely that model-generated code will do something overtly + # malicious in response to this test suite, model-generated code may act + # destructively due to a lack of model capability or alignment. + # Users are strongly encouraged to sandbox this evaluation suite so that it + # does not perform destructive actions on their host or network. + # Once you have read this disclaimer and taken appropriate precautions, + # uncomment the following line and proceed at your own risk: + exec_result = subprocess.run(["node", "test.js"], timeout=timeout, capture_output=True) + + if exec_result.stderr.decode(): + err = exec_result.stderr.decode() + result.append(f"failed: {err}") + elif exec_result.stdout.decode(): + err = exec_result.stdout.decode() + result.append(f"failed: {err}") + else: + result.append("passed") + + except TimeoutException: + result.append("timed out") + + shutil.rmtree(tmp_dir) + elif "cpp" in language_type.lower(): + import os + import shutil + + if "tmp" not in tmp_dir: + tmp_dir = os.path.join(tmp_dir, "tmp") + tmp_dir = os.path.join(tmp_dir, f"{task_id.replace('/', '-')}-{random_id}") + if not os.path.exists(tmp_dir): + os.makedirs(tmp_dir) + + os.chdir(tmp_dir) + open(f"test.cpp", 'w').write(sample["test_code"]) + if "162" in task_id: + compilation_result = subprocess.run(["/usr/bin/g++", "-std=c++11", "test.cpp", "-lcrypto", "-lssl"], + timeout=timeout, + capture_output=True) + else: + compilation_result = subprocess.run(["/usr/bin/g++", "-std=c++11", "test.cpp"], timeout=timeout, + capture_output=True) + if compilation_result.returncode != 0: + if compilation_result.stderr: + err = compilation_result.stderr.decode() + else: + err = compilation_result.stdout.decode() + result.append(f"failed: compilation error: {err}") + else: + try: + exec_result = None + with time_limit(timeout): + # WARNING + # This program exists to execute untrusted model-generated code. Although + # it is highly unlikely that model-generated code will do something overtly + # malicious in response to this test suite, model-generated code may act + # destructively due to a lack of model capability or alignment. + # Users are strongly encouraged to sandbox this evaluation suite so that it + # does not perform destructive actions on their host or network. + # Once you have read this disclaimer and taken appropriate precautions, + # uncomment the following line and proceed at your own risk: + exec_result = subprocess.run(["./a.out"], timeout=timeout, capture_output=True) + + if exec_result.returncode == 0: + result.append("passed") + else: + if exec_result.stderr: + try: + err = exec_result.stderr.decode() + except: + err = exec_result.stderr + else: + try: + err = exec_result.stdout.decode() + except: + err = exec_result.stdout + result.append(f"failed: {err}") + except TimeoutException: + result.append("timed out") + + shutil.rmtree(tmp_dir) + elif "rust" in language_type.lower(): + import os + + WD: str = os.path.dirname(os.path.abspath(__file__)) + RUST_DIR: str = os.path.join(WD, "rust") + RUST_SRC: str = os.path.join(RUST_DIR, "src") + RUST_BIN: str = os.path.join(RUST_SRC, "bin") + RUST_TMP_DIR: str = os.path.join(RUST_DIR, "tmp") + RUST_LOGS: str = os.path.join(RUST_TMP_DIR, "logs") + RUST_EXT: str = ".rs" + + # Create mandatory tmp directories + os.makedirs(RUST_TMP_DIR, exist_ok=True) + os.makedirs(RUST_LOGS, exist_ok=True) + os.makedirs(RUST_SRC, exist_ok=True) + os.makedirs(RUST_BIN, exist_ok=True) + + with tempfile.NamedTemporaryFile(dir = RUST_BIN, delete=False) as f: + #temporal file name + file_prefix = sample["task_id"].lower().replace("/", "_") + file_name:str = file_prefix +RUST_EXT + + os.rename(f.name, os.path.join(RUST_BIN, file_name)) + + # Sample to pure Rust function + rust_code: str = sample["test_code"] + + # dump the rust source code in the target temporal file + f.write(rust_code.encode('utf-8')) + + # Proceed towards Rust binaries compilation. Therefore move to Rust module root dir. + os.chdir(RUST_DIR) + + # Two possible outcomes + # Pass OR Fail compilation + log_filename: str = file_prefix + ".jsonl" + log_path: str = os.path.join(RUST_LOGS, log_filename) + cargo_check: str = "cargo check --bin " + file_prefix + " --message-format json >> " + log_path + # Compilation build status + returned_val_compilation: int + + # Overwrite file content + if os.path.exists(log_path): + if(file_size := os.path.getsize(log_path)) >= 0: + os.remove(log_path) + returned_val_compilation = os.system(cargo_check) + + else: + returned_val_compilation = os.system(cargo_check) + + # 0 means success + if returned_val_compilation == 0: + + #Execution pipeline + cargo_test: str = "cargo test --bin " +file_prefix+ " --message-format json >> " + log_path + returned_val_execution = os.system(cargo_test) + + if returned_val_execution == 0: + result.append("passed") + else: + result.append(f"failed: execution error") + + else: + result.append(f"failed: compilation error") + + + elif "java" in language_type.lower(): + assert tmp_dir is not None, "Java should be evaluated in a temporary dir." + + import os + import shutil + + if "tmp" not in tmp_dir: + tmp_dir = os.path.join(tmp_dir, "tmp") + tmp_dir = os.path.join(tmp_dir, f"{task_id.replace('/', '-')}-{random_id}") + if not os.path.exists(tmp_dir): + os.makedirs(tmp_dir) + + os.chdir(tmp_dir) + open(os.path.join(tmp_dir, "Main.java"), 'w').write(sample["test_code"]) + res = "failed: unknown error" + compile_returncode = -1 + for _ in range(5): + try: + compilation_result = subprocess.run(['javac', os.path.join(tmp_dir, "Main.java")], timeout=5, + capture_output=True) + compile_returncode = compilation_result.returncode + break + except subprocess.TimeoutExpired as e: + continue + if compile_returncode != 0: + res = "failed: compilation error" + else: + exec_result = None + try: + # WARNING + # This program exists to execute untrusted model-generated code. Although + # it is highly unlikely that model-generated code will do something overtly + # malicious in response to this test suite, model-generated code may act + # destructively due to a lack of model capability or alignment. + # Users are strongly encouraged to sandbox this evaluation suite so that it + # does not perform destructive actions on their host or network. + # Once you have read this disclaimer and taken appropriate precautions, + # uncomment the following line and proceed at your own risk: + # exec_result = subprocess.run([f'java', '-cp', tmp_dir, 'Main'], timeout=timeout, capture_output=True) + if exec_result.returncode == 0: + res = "passed" + elif exec_result.returncode == 1: + if "AssertionError" in exec_result.stderr.decode('unicode-escape'): + res = "failed: wrong answer" + else: + res = f"failed: {exec_result.stderr.decode()}" + except subprocess.TimeoutExpired as e: + res = "time out" + except BaseException as e: + res = f"failed: {e}" + result.append(res) + + shutil.rmtree(tmp_dir) + + manager = multiprocessing.Manager() + result = manager.list() + + p = multiprocessing.Process(target=unsafe_execute, args=(tmp_dir,)) + p.start() + p.join(timeout=timeout + 1) + if p.is_alive(): + p.kill() + + if not result: + result.append("timed out") + + return { + "task_id" : task_id, + "completion_id": completion_id, + "test_code" : sample["test_code"], + "prompt" : sample["prompt"], + "generation" : sample["generation"], + "result" : result[0], + "passed" : result[0] == "passed", + "finish" : -1 if "finish" not in sample else sample["finish"], + "file" : "" if "file" not in sample else sample["file"], + "output" : [] if "output" not in sample else sample["output"], + } + +# Copyright (c) OpenAI (https://openai.com) + +# Permission is hereby granted, free of charge, to any person obtaining a copy +# of this software and associated documentation files (the "Software"), to deal +# in the Software without restriction, including without limitation the rights +# to use, copy, modify, merge, publish, distribute, sublicense, and/or sell +# copies of the Software, and to permit persons to whom the Software is +# furnished to do so, subject to the following conditions: + +# The above copyright notice and this permission notice shall be included in +# all copies or substantial portions of the Software. + +# THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR +# IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY, +# FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE +# AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER +# LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM, +# OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN +# THE SOFTWARE. +# ============================================================================ +@contextlib.contextmanager +def time_limit(seconds: float): + def signal_handler(signum, frame): + raise TimeoutException("Timed out!") + + signal.setitimer(signal.ITIMER_REAL, seconds) + signal.signal(signal.SIGALRM, signal_handler) + try: + yield + finally: + signal.setitimer(signal.ITIMER_REAL, 0) + + +@contextlib.contextmanager +def swallow_io(): + stream = WriteOnlyStringIO() + with contextlib.redirect_stdout(stream): + with contextlib.redirect_stderr(stream): + with redirect_stdin(stream): + yield + + +@contextlib.contextmanager +def create_tempdir(): + with tempfile.TemporaryDirectory() as dirname: + with chdir(dirname): + yield dirname + + +class TimeoutException(Exception): + pass + + +class WriteOnlyStringIO(io.StringIO): + """ StringIO that throws an exception when it's read from """ + + def read(self, *args, **kwargs): + raise IOError + + def readline(self, *args, **kwargs): + raise IOError + + def readlines(self, *args, **kwargs): + raise IOError + + def readable(self, *args, **kwargs): + """ Returns True if the IO object can be read. """ + return False + + +class redirect_stdin(contextlib._RedirectStream): # type: ignore + _stream = 'stdin' + + +@contextlib.contextmanager +def chdir(root): + if root == ".": + yield + return + cwd = os.getcwd() + os.chdir(root) + try: + yield + except BaseException as exc: + raise exc + finally: + os.chdir(cwd) + + +def reliability_guard(maximum_memory_bytes: Optional[int] = None): + """ + This disables various destructive functions and prevents the generated code + from interfering with the test (e.g. fork bomb, killing other processes, + removing filesystem files, etc.) + + WARNING + This function is NOT a security sandbox. Untrusted code, including, model- + generated code, should not be blindly executed outside of one. See the + Codex paper for more information about OpenAI's code sandbox, and proceed + with caution. + """ + + if maximum_memory_bytes is not None: + import resource + resource.setrlimit(resource.RLIMIT_AS, (maximum_memory_bytes, maximum_memory_bytes)) + resource.setrlimit(resource.RLIMIT_DATA, (maximum_memory_bytes, maximum_memory_bytes)) + if not platform.uname().system == 'Darwin': + resource.setrlimit(resource.RLIMIT_STACK, (maximum_memory_bytes, maximum_memory_bytes)) + + faulthandler.disable() + + import builtins + builtins.exit = None + builtins.quit = None + + import os + os.environ['OMP_NUM_THREADS'] = '1' + + os.kill = None + os.system = None + os.putenv = None + os.remove = None + os.removedirs = None + os.rmdir = None + os.fchdir = None + os.setuid = None + os.fork = None + os.forkpty = None + os.killpg = None + os.rename = None + os.renames = None + os.truncate = None + os.replace = None + os.unlink = None + os.fchmod = None + os.fchown = None + os.chmod = None + os.chown = None + os.chroot = None + os.fchdir = None + os.lchflags = None + os.lchmod = None + os.lchown = None + os.getcwd = None + os.chdir = None + + import shutil + shutil.rmtree = None + shutil.move = None + shutil.chown = None + + import subprocess + subprocess.Popen = None # type: ignore + + __builtins__['help'] = None + + import sys + sys.modules['ipdb'] = None + sys.modules['joblib'] = None + sys.modules['resource'] = None + sys.modules['psutil'] = None + sys.modules['tkinter'] = None diff --git a/PyTorch/built-in/foundation/CodeGeeX2/benchmark/gather_output.py b/PyTorch/built-in/foundation/CodeGeeX2/benchmark/gather_output.py new file mode 100644 index 0000000000000000000000000000000000000000..db63986f3c078447924a935861f6a8daa2e6269f --- /dev/null +++ b/PyTorch/built-in/foundation/CodeGeeX2/benchmark/gather_output.py @@ -0,0 +1,56 @@ +import os +import sys +import fire +import glob + + +def gather_output( + output_dir: str = "./output", + output_prefix: str = None, + if_remove_rank_files: int = 0, +): + if output_prefix is None: + output_list = glob.glob(output_dir + "/*") + else: + output_list = glob.glob(os.path.join(output_dir, output_prefix + "*")) + + for output_file in output_list: + if "rank0" in output_file: + output_prefix_ = output_file.split("_rank0.jsonl")[0] + rank_files = glob.glob(output_prefix_ + "_rank*") + with open(output_prefix_ + ".jsonl", "w") as f_out: + for rank_file in rank_files: + with open(rank_file, "r") as f_in: + for line in f_in: + f_out.write(line) + if if_remove_rank_files: + os.remove(rank_file) + print(f"Removing {rank_file}...") + + if output_prefix is None: + output_list = glob.glob(output_dir + "/*") + else: + output_list = glob.glob(os.path.join(output_dir, output_prefix + "*")) + + for output_file in output_list: + if "rank" in output_file or "_unfinished" in output_file or "all" in output_file or "_result" in output_file: + continue + if "_finished" not in output_file: + continue + output_prefix_ = output_file.split("_finished.jsonl")[0] + files = [output_file, output_prefix_ + "_unfinished.jsonl"] + with open(output_prefix_ + "_all.jsonl", "w") as f_out: + for f in files: + with open(f, "r") as f_in: + for line in f_in: + f_out.write(line) + + print("Gathering finished. Saved in {}".format(output_prefix_ + "_all.jsonl")) + + +def main(): + fire.Fire(gather_output) + + +if __name__ == "__main__": + sys.exit(main()) diff --git a/PyTorch/built-in/foundation/CodeGeeX2/benchmark/humanevalx/__init__.py b/PyTorch/built-in/foundation/CodeGeeX2/benchmark/humanevalx/__init__.py new file mode 100644 index 0000000000000000000000000000000000000000..e69de29bb2d1d6434b8b29ae775ad8c2e48c5391 diff --git a/PyTorch/built-in/foundation/CodeGeeX2/benchmark/humanevalx/cpp/data/humaneval_cpp.jsonl.gz b/PyTorch/built-in/foundation/CodeGeeX2/benchmark/humanevalx/cpp/data/humaneval_cpp.jsonl.gz new file mode 100644 index 0000000000000000000000000000000000000000..e3409c83d19c3071fb5904926e6134e5d8abee18 Binary files /dev/null and b/PyTorch/built-in/foundation/CodeGeeX2/benchmark/humanevalx/cpp/data/humaneval_cpp.jsonl.gz differ diff --git a/PyTorch/built-in/foundation/CodeGeeX2/benchmark/humanevalx/cpp/evaluation/test.cpp b/PyTorch/built-in/foundation/CodeGeeX2/benchmark/humanevalx/cpp/evaluation/test.cpp new file mode 100644 index 0000000000000000000000000000000000000000..e69de29bb2d1d6434b8b29ae775ad8c2e48c5391 diff --git a/PyTorch/built-in/foundation/CodeGeeX2/benchmark/humanevalx/evaluate_humaneval_x.py b/PyTorch/built-in/foundation/CodeGeeX2/benchmark/humanevalx/evaluate_humaneval_x.py new file mode 100644 index 0000000000000000000000000000000000000000..076e0e89377192bb065bcff12e3f603b36f8111f --- /dev/null +++ b/PyTorch/built-in/foundation/CodeGeeX2/benchmark/humanevalx/evaluate_humaneval_x.py @@ -0,0 +1,239 @@ +import os +import sys +import fire +import json +import gzip +import regex +import numpy as np + +from typing import * +from tqdm.auto import tqdm +from collections import defaultdict +from concurrent.futures import ThreadPoolExecutor, as_completed + +from codegeex.benchmark.utils import read_dataset, IMPORT_HELPER +from codegeex.benchmark.metric import estimate_pass_at_k +from codegeex.benchmark.execution import check_correctness + +LANGUAGE_NAME = { + "cpp" : "CPP", + "go" : "Go", + "java" : "Java", + "js" : "JavaScript", + "python": "Python", +} + + +def process_humaneval_test(sample, problems, example_test=False): + task_id = sample["task_id"] + language = task_id.split("/")[0].lower() + + prompt = sample["prompt"] + if example_test and "example_test" in problems[task_id] and problems[task_id]["example_test"] != "": + test = problems[task_id]["example_test"] + else: + test = problems[task_id]["test"] + code = sample["generation"] + + # Pre-process for different languages + if language == "python": + code_ = [] + for line in code.split("\n"): + if (len(line.strip()) > 0 and line[0] != ' ' and line[0] != '\t'): + break + code_.append(line) + code = "\n".join(code_) + test_setup = "\n".join(IMPORT_HELPER["python"]) + "\n" + test_string = test_setup + prompt + code + "\n" + test + "\n" + elif language == "cpp": + test_set_up = "" + for s in IMPORT_HELPER["cpp"]: + if s not in prompt: + test_set_up += s + "\n" + test_string = test_set_up + "\n" + prompt + code + "\n" + test + elif language == "java": + test_string = prompt + code + "\n" + test + elif language == "js" or language == "javascript": + test_string = prompt + code + "\n" + test + elif language == "go": + import_string = problems[task_id]["import"] + prompt = prompt.replace(import_string, "") + if example_test and "example_test" in problems[task_id]: + test = problems[task_id]["example_test"] + else: + test = problems[task_id]["test"] + test_setup = problems[task_id]["test_setup"] + other_pkgs = [] + for pkg in IMPORT_HELPER["go"]: + if pkg not in test_setup: + p = pkg.split("/")[-1] + if p + "." in code: + other_pkgs.append(f"\"{pkg}\"") + if other_pkgs: + import_other_pkgs = "import (\n" + " ".join([p + "\n" for p in other_pkgs]) + ")" + test_string = test_setup + "\n" + import_other_pkgs + "\n" + prompt + code + "\n" + test + else: + test_string = test_setup + "\n" + prompt + code + "\n" + test + elif language == "rust": + main = "\nfn main(){ \n } \n" + declaration = problems[task_id]["declaration"] + test_string = main + declaration + prompt + code + test + + return test_string + + +def stream_jsonl_all(filename: str) -> Iterable[Dict]: + results = [] + if filename.endswith(".gz"): + fp = gzip.open(open(filename, "rb"), "rt") + else: + fp = open(filename, "r") + for line in fp: + if any(not x.isspace() for x in line): + results.append(json.loads(line)) + fp.close() + + return results + + +def evaluate_functional_correctness( + input_file: str = None, + tmp_dir: str = "./", + n_workers: int = 32, + timeout: float = 500.0, + problem_file: str = "../data/humaneval_python.jsonl.gz", + out_dir: str = None, + k: List[int] = [1, 10, 100], + test_groundtruth: bool = False, + example_test: bool = False, +): + if example_test: + print("Example test...") + + problems = read_dataset(problem_file, + dataset_type="humaneval") + sample_jsonl = stream_jsonl_all(input_file) + + if example_test: + suffix = "_example_test.jsonl" + else: + suffix = "_results.jsonl" + if out_dir is not None: + if not os.path.exists(out_dir): + os.makedirs(out_dir) + out_file = os.path.join(out_dir, input_file.split('/')[-1].replace(".jsonl", suffix)) + else: + out_file = os.path.join(input_file.replace(".jsonl", suffix)) + + if "/codegeex/benchmark/humaneval-x/" in input_file: + test_groundtruth = True + + if "-to-" in input_file: + translation_mode = True + else: + translation_mode = False + + with ThreadPoolExecutor(max_workers=n_workers) as executor: + + futures = [] + completion_id = Counter() + n_samples = 0 + results = defaultdict(list) + + if test_groundtruth: + print("Testing ground truth...") + for sample in tqdm(problems.values()): + task_id = sample["task_id"] + lang = task_id.split("/")[0].lower() + if lang == "javascript": + lang = "js" + tmp_dir_ = os.path.join(tmp_dir, lang, "evaluation") + sample["generation"] = sample["canonical_solution"] + sample["test_code"] = process_humaneval_test(sample, problems, example_test) + if sample["test_code"] is None: + continue + args = (task_id, sample, lang, timeout, tmp_dir_, completion_id[task_id]) + future = executor.submit(check_correctness, *args) + futures.append(future) + completion_id[task_id] += 1 + n_samples += 1 + else: + print("Reading samples...") + for sample in tqdm(sample_jsonl): + task_id = sample["task_id"] + lang = task_id.split("/")[0].lower() + if translation_mode: + task_id = sample["task_id"].split("/")[-1] + lang = regex.findall("-to-.*-", input_file)[0].split("-to-")[-1].rstrip("-") + for l in LANGUAGE_NAME: + if l in lang: + lang = l + break + task_id = f"{LANGUAGE_NAME[lang]}/{task_id}" + if lang == "javascript": + lang = "js" + tmp_dir_ = os.path.join(tmp_dir, lang, "evaluation") + sample["task_id"] = task_id + sample["test_code"] = process_humaneval_test(sample, problems, example_test) + if sample["test_code"] is None: + continue + if "completion_id" in sample: + completion_id_ = sample["completion_id"] + else: + completion_id_ = completion_id[task_id] + args = (task_id, sample, lang, timeout, tmp_dir_, completion_id_) + future = executor.submit(check_correctness, *args) + futures.append(future) + completion_id[task_id] += 1 + n_samples += 1 + + print(completion_id) + if len(completion_id) == len(problems): + evaluate_pass_at_k = True + else: + evaluate_pass_at_k = False + + print("Running test suites...") + for future in tqdm(as_completed(futures), total=len(futures)): + result = future.result() + results[result["task_id"]].append((result["completion_id"], result)) + + # Calculate pass@k. + total, correct = [], [] + for result in results.values(): + passed = [r[1]["passed"] for r in result] + total.append(len(passed)) + correct.append(sum(passed)) + total = np.array(total) + correct = np.array(correct) + if evaluate_pass_at_k: + ks = k + pass_at_k = {f"pass@{k}": estimate_pass_at_k(total, correct, k).mean() + for k in ks if (total >= k).all()} + print(pass_at_k) + else: + print("Total:", np.sum(total)) + print("Correct:", np.sum(correct)) + + print("Writing to: ", out_file) + if out_file.endswith(".gz"): + fp = gzip.GzipFile(fileobj=open(out_file, "wb"), mode="wb") + for res in results.values(): + for r in res: + fp.write((json.dumps(r[1]) + "\n").encode("utf-8")) + else: + fp = open(out_file, 'w') + for res in results.values(): + for r in res: + fp.write(json.dumps(r[1]) + "\n") + fp.close() + + print("Evaluation finished.") + + +def main(): + fire.Fire(evaluate_functional_correctness) + + +if __name__ == "__main__": + sys.exit(main()) diff --git a/PyTorch/built-in/foundation/CodeGeeX2/benchmark/humanevalx/generate_humaneval_x.py b/PyTorch/built-in/foundation/CodeGeeX2/benchmark/humanevalx/generate_humaneval_x.py new file mode 100644 index 0000000000000000000000000000000000000000..75eee09d0e23e10c6537b20b6a94563579281ef2 --- /dev/null +++ b/PyTorch/built-in/foundation/CodeGeeX2/benchmark/humanevalx/generate_humaneval_x.py @@ -0,0 +1,441 @@ +import argparse +import logging +import os +import random +import socket +import time +from typing import * + +import torch +import zmq + +from codegeex.benchmark.utils import read_dataset, process_extra_prompt +from codegeex.megatron import get_args +from codegeex.megatron.inference import run_generation_distributed, model_provider +from codegeex.megatron.initialize import initialize_megatron +from codegeex.quantization import quantize + +logging.getLogger("torch").setLevel(logging.WARNING) + + +def add_code_generation_args(parser): + """Code generation arguments.""" + group = parser.add_argument_group(title="code generation") + + group.add_argument( + "--hostfile", + type=str, + default="./hostfile", + ) + group.add_argument( + "--channel-ip", + type=str, + default=None, + help="IP for ZeroMQ channel", + ) + group.add_argument( + "--channel-port", + type=int, + default=5555, + help="Port for ZeroMQ channel", + ) + group.add_argument( + "--master-port", + type=int, + default=6666, + help="Port for torch distributed", + ) + group.add_argument( + "--temperature", + type=float, + default=1.0, + help="Sampling temperature.", + ) + group.add_argument( + "--greedy", + action="store_true", + default=False, + help="Use greedy sampling.", + ) + group.add_argument( + "--top-p", + type=float, + default=0.0, + help="Top p sampling.", + ) + group.add_argument( + "--top-k", + type=int, + default=0, + help="Top k sampling.", + ) + group.add_argument( + "--out-seq-length", + type=int, + default=1024, + help="Size of the output generated text.", + ) + group.add_argument( + "--input-path", + type=str, + default="./benchmark/humaneval/HumanEval.jsonl", + help="Get input path", + ) + group.add_argument( + "--num-samples", + type=int, + default=0, + help="Number of samples to generate", + ) + group.add_argument( + "--recompute", + action="store_true", + help="During generation recompute all attention " + "instead of using previously computed keys/values.", + ) + group.add_argument( + "--load-deepspeed", + action="store_true", + help="Load DeepSpeed checkpoint", + ) + group.add_argument( + "--ws-encoding-start-id", + type=int, + default=None, + help="Start id for whitespace encoding", + ) + group.add_argument( + "--ws-encoding-length", + type=int, + default=None, + help="Length of whitespace encoding", + ) + group.add_argument( + "--dataset", + type=str, + default="humaneval", + ) + group.add_argument( + "--samples-per-problem", + type=int, + default=200, + help="Number of samples to generate for each problem", + ) + group.add_argument( + "--output-prefix", + type=str, + default="./output/humaneval", + help="Prefix for output files", + ) + group.add_argument( + "--gen-node-rank", + type=int, + default=None, + ) + group.add_argument( + "--gen-node-world-size", + type=int, + default=None, + ) + group.add_argument( + "--gen-world-size", + type=int, + default=1, + help="Number of machines to use for generation", + ) + group.add_argument( + "--gen-rank", + type=int, + default=0, + help="Machine rank for human eval generation", + ) + group.add_argument( + "--extra-prompt", + type=str, + default=None, + help="Extra prompt to use for human eval generation", + ) + group.add_argument( + "--verbose-interval", + type=int, + default=100, + ) + group.add_argument( + "--problem-split", + type=str, + default="test", + ) + group.add_argument( + "--prompt-type", + type=str, + default="notag", + ) + group.add_argument( + "--num-devices-per-node", + type=int, + default=None, + ) + group.add_argument( + "--return-scores", + action="store_true", + ) + group.add_argument( + '--language-type', + default=None, + help='Identify the type of programming language to generate', + ) + group.add_argument( + '--bad-ids', + nargs="*", + type=int, + default=None, + help='Specify bad ids that will not be used', + ) + group.add_argument( + "--quantize", + action="store_true", + ) + + return parser + + +def main(node_rank: int, local_rank: int, master_port: int, num_devices: int): + """Main program.""" + os.environ["WORLD_SIZE"] = str(num_devices) + os.environ["RANK"] = str(local_rank) + os.environ["MASTER_ADDR"] = "0.0.0.0" + os.environ["MASTER_PORT"] = f"{master_port}" + + initialize_megatron( + extra_args_provider=add_code_generation_args, + args_defaults={ + "tokenizer_type": "GPT2BPETokenizer", + "no_load_rng" : True, + "no_load_optim" : True, + }, + ) + + # set_random_seed(node_rank * num_devices + local_rank) + args = get_args() + if args.num_layers_per_virtual_pipeline_stage is not None: + print("Interleaved pipeline schedule is not yet supported for text generation.") + exit() + + world_size = args.gen_node_world_size * num_devices + args.gen_rank = num_devices * node_rank + local_rank + args.gen_world_size = world_size + print(f"Generating on rank {args.gen_rank} of {args.gen_world_size}") + + # Set up model and load checkpoint. + state_dict = torch.load(args.load, map_location="cpu") + state_dict = state_dict["module"] + + print("Building CodeGeeX model ...") + model = model_provider() + model.load_state_dict(state_dict) + model.eval() + if args.fp16 and args.ln_fp16: + model.half() + if args.quantize: + model = quantize(model, weight_bit_width=8, backend="megatron") + model.cuda() + + # Generate samples. + run_generation_distributed(model) + + print(f"(gen_rank={args.gen_rank}, rank={local_rank}) finished, waiting ...") + torch.distributed.barrier() + + +def server(): + print(f"[ server ] starting ...", flush=True) + parser = argparse.ArgumentParser() + parser.add_argument( + "--channel-ip", + type=str, + default=None, + help="IP for ZeroMQ channel", + ) + parser.add_argument( + "--channel-port", + type=int, + default=5555, + help="Port for ZeroMQ channel", + ) + parser.add_argument( + "--master-port", + type=int, + default=6666, + help="Port for torch distributed", + ) + parser.add_argument( + "--samples-per-problem", + type=int, + default=200, + help="Number of samples to generate for each problem", + ) + parser.add_argument( + "--gen-node-world-size", + type=int, + default=1, + help="Number of machines to use for generation", + ) + parser.add_argument( + "--input-path", + type=str, + default="", + help="Get input path", + ) + parser.add_argument( + "--problem-split", + type=str, + default="test", + ) + parser.add_argument( + "--micro-batch-size", + type=int, + default=1, + ) + parser.add_argument( + '--language-type', + default=None, + help='Identify the type of programming language to generate', + ) + + args = parser.parse_known_args()[0] + entries = read_dataset(args.input_path, dataset_type="humaneval") + + assert args.samples_per_problem % args.micro_batch_size == 0, "samples_per_problem should be divisible by micro_batch_size" + + for entry in entries.values(): + entry["prompt"] = process_extra_prompt(entry["prompt"], args.language_type) + + res = [] + for entry in entries.values(): + res.extend([entry] * (args.samples_per_problem // args.micro_batch_size)) + random.shuffle(res) + all_entries = res + + # setup zeromq channel + print(f"[ server ] starting up on port {args.channel_port}", flush=True) + context = zmq.Context() + print(f"[ server ] creating socket", flush=True) + socket = context.socket(zmq.REP) + print(f"[ server ] binding to port {args.channel_port}", flush=True) + socket.bind(f"tcp://*:{args.channel_port}") + + print( + f"[ server ] loaded {len(entries)} entries, generated {len(all_entries)} samples", + flush=True, + ) + + remaining_entries = all_entries.copy() + running_workers = args.gen_node_world_size * torch.cuda.device_count() + num_finished = 0 + + print(f"[ server ] listening for requests ...", flush=True) + start_time = time.perf_counter() + while True: + # Wait for next request from client + msg = socket.recv_json() + rank = msg["rank"] + action = msg["action"] + + if action == "pull": + if len(remaining_entries) == 0: + print(f"[ server ] Shutting down worker {rank}", flush=True) + socket.send_json({"task_id": None}) + running_workers -= 1 + if running_workers == 0: + print(f"[ server ] All workers finished", flush=True) + break + else: + entry = remaining_entries.pop() + time_elapsed = time.perf_counter() - start_time + print(f"[ server ] Sending entry {entry['task_id']} to worker {rank}", flush=True) + remaining = ( + len(remaining_entries) + / (len(all_entries) - len(remaining_entries)) + * time_elapsed + ) + time_per_sampple = 0.0 if num_finished == 0 else time_elapsed / num_finished / args.micro_batch_size + print( + f"[ server ] total {len(all_entries)}, assigned {len(all_entries) - len(remaining_entries)}, " + f"finished {num_finished}, " + f"elapsed {time_elapsed:.4f}", + f"speed {time_per_sampple:.4f}s/sample", + f"remaining {remaining:.4f}", + flush=True, + ) + socket.send_json({"task_id": entry}) + else: + if action == "success": + print(f"[ server ] {msg['task_id']} is finished", flush=True) + socket.send_json({"pong": 1}) + else: + print(f"[ server ] {msg['task_id']} is not finished", flush=True) + remaining_entries.append(msg['task_id']) + socket.send_json({"pong": 1}) + break + + num_finished += 1 + + +if __name__ == "__main__": + torch.multiprocessing.set_start_method("spawn") + parser = argparse.ArgumentParser() + parser.add_argument( + "--hostfile", + type=str, + default="./hostfile", + ) + parser.add_argument( + "--master-port", + type=int, + default=5666, + ) + parser.add_argument( + "--seed", + type=int, + default=42, + ) + args = parser.parse_known_args()[0] + + print("start method: " + torch.multiprocessing.get_start_method()) + + processes = [] + num_devices = torch.cuda.device_count() + + hosts = open(args.hostfile, "r").readlines() + hosts = [host.strip() for host in hosts] + master_port = args.master_port + + node_rank = None + for i in range(len(hosts)): + if hosts[i] == socket.gethostbyname(socket.gethostname()): + node_rank = i + break + assert ( + node_rank is not None + ), f"Could not find hostname ({socket.gethostbyname(socket.gethostname())}) in hostlist" + + # launch server + if socket.gethostbyname(socket.gethostname()) == hosts[0]: + server_process = torch.multiprocessing.Process(target=server) + print(f"Launching server ...") + server_process.start() + processes.append(server_process) + + for i in range(num_devices): + local_rank = i + print(f"launching local rank {i}") + + p = torch.multiprocessing.Process( + target=main, + args=(node_rank, local_rank, master_port, num_devices), + ) + p.start() + processes.append(p) + + for p in processes: + p.join() diff --git a/PyTorch/built-in/foundation/CodeGeeX2/benchmark/humanevalx/go/go.mod b/PyTorch/built-in/foundation/CodeGeeX2/benchmark/humanevalx/go/go.mod new file mode 100644 index 0000000000000000000000000000000000000000..80a6322a6cfec260fae97c9d973103a136236b07 --- /dev/null +++ b/PyTorch/built-in/foundation/CodeGeeX2/benchmark/humanevalx/go/go.mod @@ -0,0 +1,14 @@ +module humanEval + +go 1.18 + +require ( + github.com/go-openapi/inflect v0.19.0 + github.com/stretchr/testify v1.8.0 +) + +require ( + github.com/davecgh/go-spew v1.1.1 // indirect + github.com/pmezard/go-difflib v1.0.0 // indirect + gopkg.in/yaml.v3 v3.0.1 // indirect +) diff --git a/PyTorch/built-in/foundation/CodeGeeX2/benchmark/humanevalx/go/go.sum b/PyTorch/built-in/foundation/CodeGeeX2/benchmark/humanevalx/go/go.sum new file mode 100644 index 0000000000000000000000000000000000000000..d88557079e79dba9786aee16537cd1c35c2fbc1e --- /dev/null +++ b/PyTorch/built-in/foundation/CodeGeeX2/benchmark/humanevalx/go/go.sum @@ -0,0 +1,17 @@ +github.com/davecgh/go-spew v1.1.0/go.mod h1:J7Y8YcW2NihsgmVo/mv3lAwl/skON4iLHjSsI+c5H38= +github.com/davecgh/go-spew v1.1.1 h1:vj9j/u1bqnvCEfJOwUhtlOARqs3+rkHYY13jYWTU97c= +github.com/davecgh/go-spew v1.1.1/go.mod h1:J7Y8YcW2NihsgmVo/mv3lAwl/skON4iLHjSsI+c5H38= +github.com/go-openapi/inflect v0.19.0 h1:9jCH9scKIbHeV9m12SmPilScz6krDxKRasNNSNPXu/4= +github.com/go-openapi/inflect v0.19.0/go.mod h1:lHpZVlpIQqLyKwJ4N+YSc9hchQy/i12fJykb83CRBH4= +github.com/pmezard/go-difflib v1.0.0 h1:4DBwDE0NGyQoBHbLQYPwSUPoCMWR5BEzIk/f1lZbAQM= +github.com/pmezard/go-difflib v1.0.0/go.mod h1:iKH77koFhYxTK1pcRnkKkqfTogsbg7gZNVY4sRDYZ/4= +github.com/stretchr/objx v0.1.0/go.mod h1:HFkY916IF+rwdDfMAkV7OtwuqBVzrE8GR6GFx+wExME= +github.com/stretchr/objx v0.4.0/go.mod h1:YvHI0jy2hoMjB+UWwv71VJQ9isScKT/TqJzVSSt89Yw= +github.com/stretchr/testify v1.7.1/go.mod h1:6Fq8oRcR53rry900zMqJjRRixrwX3KX962/h/Wwjteg= +github.com/stretchr/testify v1.8.0 h1:pSgiaMZlXftHpm5L7V1+rVB+AZJydKsMxsQBIJw4PKk= +github.com/stretchr/testify v1.8.0/go.mod h1:yNjHg4UonilssWZ8iaSj1OCr/vHnekPRkoO+kdMU+MU= +gopkg.in/check.v1 v0.0.0-20161208181325-20d25e280405 h1:yhCVgyC4o1eVCa2tZl7eS0r+SDo693bJlVdllGtEeKM= +gopkg.in/check.v1 v0.0.0-20161208181325-20d25e280405/go.mod h1:Co6ibVJAznAaIkqp8huTwlJQCZ016jof/cbN4VW5Yz0= +gopkg.in/yaml.v3 v3.0.0-20200313102051-9f266ea9e77c/go.mod h1:K4uyk7z7BCEPqu6E+C64Yfv1cQ7kz7rIZviUmN+EgEM= +gopkg.in/yaml.v3 v3.0.1 h1:fxVm/GzAzEWqLHuvctI91KS9hhNmmWOoWu0XTYJS7CA= +gopkg.in/yaml.v3 v3.0.1/go.mod h1:K4uyk7z7BCEPqu6E+C64Yfv1cQ7kz7rIZviUmN+EgEM= diff --git a/PyTorch/built-in/foundation/CodeGeeX2/benchmark/humanevalx/go/vendor.tar.gz b/PyTorch/built-in/foundation/CodeGeeX2/benchmark/humanevalx/go/vendor.tar.gz new file mode 100644 index 0000000000000000000000000000000000000000..6d6c649e558abd01c23488bb906c2d4f3e805e10 Binary files /dev/null and b/PyTorch/built-in/foundation/CodeGeeX2/benchmark/humanevalx/go/vendor.tar.gz differ diff --git a/PyTorch/built-in/foundation/CodeGeeX2/benchmark/humanevalx/humanevalx_cpp.jsonl.gz b/PyTorch/built-in/foundation/CodeGeeX2/benchmark/humanevalx/humanevalx_cpp.jsonl.gz new file mode 100644 index 0000000000000000000000000000000000000000..e3409c83d19c3071fb5904926e6134e5d8abee18 Binary files /dev/null and b/PyTorch/built-in/foundation/CodeGeeX2/benchmark/humanevalx/humanevalx_cpp.jsonl.gz differ diff --git a/PyTorch/built-in/foundation/CodeGeeX2/benchmark/humanevalx/humanevalx_go.jsonl.gz b/PyTorch/built-in/foundation/CodeGeeX2/benchmark/humanevalx/humanevalx_go.jsonl.gz new file mode 100644 index 0000000000000000000000000000000000000000..ef8a94062188ab5a4c6fdce2d24769a90698c63b Binary files /dev/null and b/PyTorch/built-in/foundation/CodeGeeX2/benchmark/humanevalx/humanevalx_go.jsonl.gz differ diff --git a/PyTorch/built-in/foundation/CodeGeeX2/benchmark/humanevalx/humanevalx_java.jsonl.gz b/PyTorch/built-in/foundation/CodeGeeX2/benchmark/humanevalx/humanevalx_java.jsonl.gz new file mode 100644 index 0000000000000000000000000000000000000000..f9b1d0a477feac75ccce1a9d6f765a600aceb8ba Binary files /dev/null and b/PyTorch/built-in/foundation/CodeGeeX2/benchmark/humanevalx/humanevalx_java.jsonl.gz differ diff --git a/PyTorch/built-in/foundation/CodeGeeX2/benchmark/humanevalx/humanevalx_js.jsonl.gz b/PyTorch/built-in/foundation/CodeGeeX2/benchmark/humanevalx/humanevalx_js.jsonl.gz new file mode 100644 index 0000000000000000000000000000000000000000..234387d1a4f0105c763de3f514771813446685f8 Binary files /dev/null and b/PyTorch/built-in/foundation/CodeGeeX2/benchmark/humanevalx/humanevalx_js.jsonl.gz differ diff --git a/PyTorch/built-in/foundation/CodeGeeX2/benchmark/humanevalx/humanevalx_python.jsonl.gz b/PyTorch/built-in/foundation/CodeGeeX2/benchmark/humanevalx/humanevalx_python.jsonl.gz new file mode 100644 index 0000000000000000000000000000000000000000..f34f01385dba3a2b4458b09b3cb5f780cf18a200 Binary files /dev/null and b/PyTorch/built-in/foundation/CodeGeeX2/benchmark/humanevalx/humanevalx_python.jsonl.gz differ diff --git a/PyTorch/built-in/foundation/CodeGeeX2/benchmark/humanevalx/humanevalx_rust.jsonl.gz b/PyTorch/built-in/foundation/CodeGeeX2/benchmark/humanevalx/humanevalx_rust.jsonl.gz new file mode 100644 index 0000000000000000000000000000000000000000..634c21c3c83b210648fb2d624ee7409d2fd91577 Binary files /dev/null and b/PyTorch/built-in/foundation/CodeGeeX2/benchmark/humanevalx/humanevalx_rust.jsonl.gz differ diff --git a/PyTorch/built-in/foundation/CodeGeeX2/benchmark/humanevalx/rust/Cargo.lock b/PyTorch/built-in/foundation/CodeGeeX2/benchmark/humanevalx/rust/Cargo.lock new file mode 100644 index 0000000000000000000000000000000000000000..e2e06556aa4099da3e074da1e58c48df846a8aee --- /dev/null +++ b/PyTorch/built-in/foundation/CodeGeeX2/benchmark/humanevalx/rust/Cargo.lock @@ -0,0 +1,121 @@ +# This file is automatically @generated by Cargo. +# It is not intended for manual editing. +version = 3 + +[[package]] +name = "aho-corasick" +version = "0.7.20" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "cc936419f96fa211c1b9166887b38e5e40b19958e5b895be7c1f93adec7071ac" +dependencies = [ + "memchr", +] + +[[package]] +name = "fuchsia-cprng" +version = "0.1.1" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "a06f77d526c1a601b7c4cdd98f54b5eaabffc14d5f2f0296febdc7f357c6d3ba" + +[[package]] +name = "libc" +version = "0.2.139" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "201de327520df007757c1f0adce6e827fe8562fbc28bfd9c15571c66ca1f5f79" + +[[package]] +name = "md5" +version = "0.7.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "490cc448043f947bae3cbee9c203358d62dbee0db12107a74be5c30ccfd09771" + +[[package]] +name = "memchr" +version = "2.5.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "2dffe52ecf27772e601905b7522cb4ef790d2cc203488bbd0e2fe85fcb74566d" + +[[package]] +name = "rand" +version = "0.4.6" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "552840b97013b1a26992c11eac34bdd778e464601a4c2054b5f0bff7c6761293" +dependencies = [ + "fuchsia-cprng", + "libc", + "rand_core 0.3.1", + "rdrand", + "winapi", +] + +[[package]] +name = "rand_core" +version = "0.3.1" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "7a6fdeb83b075e8266dcc8762c22776f6877a63111121f5f8c7411e5be7eed4b" +dependencies = [ + "rand_core 0.4.2", +] + +[[package]] +name = "rand_core" +version = "0.4.2" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "9c33a3c44ca05fa6f1807d8e6743f3824e8509beca625669633be0acbdf509dc" + +[[package]] +name = "rdrand" +version = "0.4.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "678054eb77286b51581ba43620cc911abf02758c91f93f479767aed0f90458b2" +dependencies = [ + "rand_core 0.3.1", +] + +[[package]] +name = "regex" +version = "1.7.1" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "48aaa5748ba571fb95cd2c85c09f629215d3a6ece942baa100950af03a34f733" +dependencies = [ + "aho-corasick", + "memchr", + "regex-syntax", +] + +[[package]] +name = "regex-syntax" +version = "0.6.28" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "456c603be3e8d448b072f410900c09faf164fbce2d480456f50eea6e25f9c848" + +[[package]] +name = "rust" +version = "0.1.0" +dependencies = [ + "md5", + "rand", + "regex", +] + +[[package]] +name = "winapi" +version = "0.3.9" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "5c839a674fcd7a98952e593242ea400abe93992746761e38641405d28b00f419" +dependencies = [ + "winapi-i686-pc-windows-gnu", + "winapi-x86_64-pc-windows-gnu", +] + +[[package]] +name = "winapi-i686-pc-windows-gnu" +version = "0.4.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "ac3b87c63620426dd9b991e5ce0329eff545bccbbb34f3be09ff6fb6ab51b7b6" + +[[package]] +name = "winapi-x86_64-pc-windows-gnu" +version = "0.4.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "712e227841d057c1ee1cd2fb22fa7e5a5461ae8e48fa2ca79ec42cfc1931183f" diff --git a/PyTorch/built-in/foundation/CodeGeeX2/benchmark/humanevalx/rust/Cargo.toml b/PyTorch/built-in/foundation/CodeGeeX2/benchmark/humanevalx/rust/Cargo.toml new file mode 100644 index 0000000000000000000000000000000000000000..1ed0c3a5de92995e023a3edaa290751522c9de31 --- /dev/null +++ b/PyTorch/built-in/foundation/CodeGeeX2/benchmark/humanevalx/rust/Cargo.toml @@ -0,0 +1,12 @@ +[package] +name = "rust" +version = "0.1.0" +edition = "2021" + +# See more keys and their definitions at https://doc.rust-lang.org/cargo/reference/manifest.html + +[dependencies] +rand = "0.4" +regex = "1" +md5 = "0.7.0" + diff --git a/PyTorch/built-in/foundation/CodeGeeX2/benchmark/humanevalx/translate_humaneval_x.py b/PyTorch/built-in/foundation/CodeGeeX2/benchmark/humanevalx/translate_humaneval_x.py new file mode 100644 index 0000000000000000000000000000000000000000..cab2888a3cb791f43e7cbb0f1d2cd72dd15307a8 --- /dev/null +++ b/PyTorch/built-in/foundation/CodeGeeX2/benchmark/humanevalx/translate_humaneval_x.py @@ -0,0 +1,480 @@ +import os +import zmq +import time +import torch +import random +import socket +import logging +import argparse + +from typing import * +from codegeex.benchmark.utils import read_translation_dataset +from codegeex.megatron import get_args +from codegeex.megatron.inference import run_generation_distributed, model_provider +from codegeex.megatron.initialize import initialize_megatron + +logging.getLogger("torch").setLevel(logging.WARNING) + + +def add_code_generate_args(parser): + """Code generation arguments.""" + group = parser.add_argument_group(title="code generation") + + group.add_argument( + "--hostfile", + type=str, + default="./hostfile", + ) + group.add_argument( + "--channel-ip", + type=str, + default=None, + help="IP for ZeroMQ channel", + ) + group.add_argument( + "--channel-port", + type=int, + default=5555, + help="Port for ZeroMQ channel", + ) + group.add_argument( + "--master-port", + type=int, + default=5666, + ) + group.add_argument( + "--temperature", + type=float, + default=1.0, + help="Sampling temperature.", + ) + group.add_argument( + "--greedy", + action="store_true", + default=False, + help="Use greedy sampling.", + ) + group.add_argument( + "--top-p", + type=float, + default=0.0, + help="Top p sampling.", + ) + group.add_argument( + "--top-k", + type=int, + default=0, + help="Top k sampling.", + ) + group.add_argument( + "--out-seq-length", + type=int, + default=1024, + help="Size of the output generated text.", + ) + group.add_argument( + "--input-path", + type=str, + default="./benchmark/humaneval/HumanEval.jsonl", + help="Get input path", + ) + group.add_argument( + "--num-samples", + type=int, + default=0, + help="Number of samples to generate", + ) + group.add_argument( + "--recompute", + action="store_true", + help="During generation recompute all attention " + "instead of using previously computed keys/values.", + ) + group.add_argument( + "--load-deepspeed", + action="store_true", + help="Load DeepSpeed checkpoint", + ) + group.add_argument( + "--ws-encoding-start-id", + type=int, + default=None, + help="Start id for whitespace encoding", + ) + group.add_argument( + "--ws-encoding-length", + type=int, + default=None, + help="Length of whitespace encoding", + ) + group.add_argument( + "--dataset", + type=str, + default="humaneval", + ) + group.add_argument( + "--samples-per-problem", + type=int, + default=200, + help="Number of samples to generate for each problem", + ) + group.add_argument( + "--output-prefix", + type=str, + default="./output/humaneval", + help="Prefix for output files", + ) + group.add_argument( + "--gen-node-rank", + type=int, + default=None, + ) + group.add_argument( + "--gen-node-world-size", + type=int, + default=None, + ) + group.add_argument( + "--gen-world-size", + type=int, + default=1, + help="Number of machines to use for generation", + ) + group.add_argument( + "--gen-rank", + type=int, + default=0, + help="Machine rank for human eval generation", + ) + group.add_argument( + "--extra-prompt", + type=str, + default=None, + help="Extra prompt to use for human eval generation", + ) + group.add_argument( + "--verbose-interval", + type=int, + default=100, + ) + group.add_argument( + "--problem-split", + type=str, + default="test", + ) + group.add_argument( + "--prompt-type", + type=str, + default="notag", + ) + group.add_argument( + "--num-devices-per-node", + type=int, + default=None, + ) + group.add_argument( + "--return-scores", + action="store_true", + ) + group.add_argument( + "--free-guidance", + action="store_true", + ) + group.add_argument( + "--guide-temp", + type=float, + default=1.5, + ) + group.add_argument( + "--attention-upweight", + type=float, + default=None, + ) + group.add_argument( + '--bad-ids', + nargs="*", + type=int, + default=None, + help='Identify the type of programming language to generate', + ) + group.add_argument( + "--src-path", + type=str, + default="", + help="Get source path", + ) + group.add_argument( + "--tgt-path", + type=str, + default="", + help="Get target path", + ) + group.add_argument( + '--language-src-type', + type=str, + default=None, + help='Identify the type of programming language', + ) + group.add_argument( + '--language-tgt-type', + type=str, + default=None, + help='Identify the type of programming language to translate', + ) + + return parser + + +def main(node_rank: int, local_rank: int, master_port: int, num_devices: int): + """Main program.""" + os.environ["WORLD_SIZE"] = str(num_devices) + os.environ["RANK"] = str(local_rank) + os.environ["MASTER_ADDR"] = "0.0.0.0" + os.environ["MASTER_PORT"] = f"{master_port}" + + initialize_megatron( + extra_args_provider=add_code_generate_args, + args_defaults={ + "tokenizer_type": "GPT2BPETokenizer", + "no_load_rng" : True, + "no_load_optim" : True, + }, + ) + + # set_random_seed(node_rank * num_devices + local_rank) + args = get_args() + if args.num_layers_per_virtual_pipeline_stage is not None: + print("Interleaved pipeline schedule is not yet supported for text generation.") + exit() + + world_size = args.gen_node_world_size * num_devices + args.gen_rank = num_devices * node_rank + local_rank + args.gen_world_size = world_size + print(f"Generating on rank {args.gen_rank} of {args.gen_world_size}") + + # Set up model and load checkpoint. + state_dict = torch.load(args.load, map_location="cpu") + state_dict = state_dict["module"] + + print("Building CodeGeeX model ...") + model = model_provider() + model.load_state_dict(state_dict) + model.eval() + if args.fp16 and args.ln_fp16: + model.half() + model.cuda() + + # Generate samples. + run_generation_distributed(model) + + print(f"(gen_rank={args.gen_rank}, rank={local_rank}) finished, waiting ...") + torch.distributed.barrier() + + +def server(): + print(f"[ server ] starting ...", flush=True) + parser = argparse.ArgumentParser() + parser.add_argument( + "--channel-ip", + type=str, + default=None, + help="IP for ZeroMQ channel", + ) + parser.add_argument( + "--channel-port", + type=int, + default=5555, + help="Port for ZeroMQ channel", + ) + parser.add_argument( + "--master-port", + type=int, + default=6666, + help="Port for torch distributed", + ) + parser.add_argument( + "--samples-per-problem", + type=int, + default=200, + help="Number of samples to generate for each problem", + ) + parser.add_argument( + "--gen-node-world-size", + type=int, + default=1, + help="Number of machines to use for generation", + ) + parser.add_argument( + "--src-path", + type=str, + default="", + help="Get source path", + ) + parser.add_argument( + "--tgt-path", + type=str, + default="", + help="Get target path", + ) + parser.add_argument( + "--problem-split", + type=str, + default="test", + ) + parser.add_argument( + "--micro-batch-size", + type=int, + default=1, + ) + parser.add_argument( + '--language-src-type', + type=str, + default=None, + help='Identify the type of programming language', + ) + parser.add_argument( + '--language-tgt-type', + type=str, + default=None, + help='Identify the type of programming language to translate', + ) + + args = parser.parse_known_args()[0] + + entries = read_translation_dataset(args.src_path, + args.tgt_path, + lang_src=args.language_src_type, + lang_tgt=args.language_tgt_type, + dataset_type="humaneval") + + assert args.samples_per_problem % args.micro_batch_size == 0, "samples_per_problem should be divisible by micro_batch_size" + + res = [] + for entry in entries.values(): + res.extend([entry] * (args.samples_per_problem // args.micro_batch_size)) + random.shuffle(res) + all_entries = res + + # setup zeromq channel + print(f"[ server ] starting up on port {args.channel_port}", flush=True) + context = zmq.Context() + print(f"[ server ] creating socket", flush=True) + socket = context.socket(zmq.REP) + print(f"[ server ] binding to port {args.channel_port}", flush=True) + socket.bind(f"tcp://*:{args.channel_port}") + + print( + f"[ server ] loaded {len(entries)} entries, generated {len(all_entries)} samples", + flush=True, + ) + + remaining_entries = all_entries.copy() + running_workers = args.gen_node_world_size * torch.cuda.device_count() + num_finished = 0 + + print(f"[ server ] listening for requests ...", flush=True) + start_time = time.perf_counter() + while True: + # Wait for next request from client + msg = socket.recv_json() + rank = msg["rank"] + action = msg["action"] + + if action == "pull": + if len(remaining_entries) == 0: + print(f"[ server ] Shutting down worker {rank}", flush=True) + socket.send_json({"task_id": None}) + running_workers -= 1 + if running_workers == 0: + print(f"[ server ] All workers finished", flush=True) + break + else: + entry = remaining_entries.pop() + time_elapsed = time.perf_counter() - start_time + print(f"[ server ] Sending entry {entry['task_id']} to worker {rank}", flush=True) + remaining = ( + len(remaining_entries) + / (len(all_entries) - len(remaining_entries)) + * time_elapsed + ) + time_per_sampple = 0.0 if num_finished == 0 else time_elapsed / num_finished / args.micro_batch_size + print( + f"[ server ] total {len(all_entries)}, assigned {len(all_entries) - len(remaining_entries)}, " + f"finished {num_finished}, " + f"elapsed {time_elapsed:.4f}", + f"speed {time_per_sampple:.4f}s/sample", + f"remaining {remaining:.4f}", + flush=True, + ) + socket.send_json({"task_id": entry}) + else: + if action == "success": + print(f"[ server ] {msg['task_id']} is finished", flush=True) + socket.send_json({"pong": 1}) + else: + print(f"[ server ] {msg['task_id']} is not finished", flush=True) + remaining_entries.append(msg['task_id']) + socket.send_json({"pong": 1}) + break + + num_finished += 1 + + +if __name__ == "__main__": + torch.multiprocessing.set_start_method("spawn") + parser = argparse.ArgumentParser() + parser.add_argument( + "--hostfile", + type=str, + default="./hostfile", + ) + parser.add_argument( + "--master-port", + type=int, + default=5666, + ) + parser.add_argument( + "--seed", + type=int, + default=42, + ) + args = parser.parse_known_args()[0] + + print("start method: " + torch.multiprocessing.get_start_method()) + + processes = [] + num_devices = torch.cuda.device_count() + + hosts = open(args.hostfile, "r").readlines() + hosts = [host.strip() for host in hosts] + master_port = args.master_port + + node_rank = None + for i in range(len(hosts)): + if hosts[i] == socket.gethostbyname(socket.gethostname()): + node_rank = i + break + assert ( + node_rank is not None + ), f"Could not find hostname ({socket.gethostbyname(socket.gethostname())}) in hostlist" + + # launch server + if socket.gethostbyname(socket.gethostname()) == hosts[0]: + server_process = torch.multiprocessing.Process(target=server) + print(f"Launching server ...") + server_process.start() + processes.append(server_process) + + for i in range(num_devices): + local_rank = i + print(f"launching local rank {i}") + + p = torch.multiprocessing.Process( + target=main, + args=(node_rank, local_rank, master_port, num_devices), + ) + p.start() + processes.append(p) + + for p in processes: + p.join() diff --git a/PyTorch/built-in/foundation/CodeGeeX2/benchmark/inspect_result.py b/PyTorch/built-in/foundation/CodeGeeX2/benchmark/inspect_result.py new file mode 100644 index 0000000000000000000000000000000000000000..52adf995cee2698f8786dad2d8a6e29c3e2a5834 --- /dev/null +++ b/PyTorch/built-in/foundation/CodeGeeX2/benchmark/inspect_result.py @@ -0,0 +1,286 @@ +import os +import sys +import fire +import json +import glob +import numpy as np +import pandas as pd + +from collections import defaultdict +from codegeex.benchmark.metric import estimate_pass_at_k + +error_types = { + "python" : [ + "accepted", + "assertion error", + "undefined error", + "runtime error", + "syntax error", + "timeout error", + "type error", + ], + "java" : [ + "accepted", + "compilation error", + "assertion error", + "timeout error", + "index error", + "class cast error", + "stack overflow", + "null pointer error", + "unsupported operation", + "number format error", + "no such element", + "illegal argument", + "out of memory", + "arithmetic error", + "others", + ], + "cpp" : [ + "accepted", + "compilation error", + "assertion error", + "range error", + "invalid argument", + "pointer error", + "out of memory", + "package error", + "others", + ], + "javascript": [ + "accepted", + "assertion error", + "undefined error", + "runtime error", + "syntax error", + "timeout error", + "range error", + "type error", + ], + "go" : [ + "accepted", + "assertion error", + "undefined error", + "runtime error", + "syntax error", + "timeout error", + "type error", + "notused error", + ], +} + + +def inspect_result( + input_dir: str = None, + input_file: str = None, + output_dir: str = None, + pass_at_k_outpath: str = None, +): + if input_dir is not None: + input_files = glob.glob(input_dir + "/*_results.jsonl") + else: + input_files = [input_file] + + if output_dir is not None: + if not os.path.exists(output_dir): + os.makedirs(output_dir) + else: + output_dir = "/" + + pass_at_k_outs = [] + for input_file in input_files: + result_stats = defaultdict(dict) + df = None + incompleted = False + with open(input_file, "r") as f: + for i, line in enumerate(f): + obj = json.loads(line) + task_id = obj["task_id"] + language_type = task_id.split("/")[0].lower() + if language_type not in error_types: + if language_type == "humaneval": + language_type = "python" + elif language_type == "C++": + language_type = "cpp" + else: + incompleted = True + break + if task_id not in result_stats.keys(): + default_stats = {} + default_stats["task_id"] = task_id + default_stats["n_sample"] = 0 + for k in error_types[language_type]: + default_stats[k] = 0 + + result_stats[task_id] = default_stats.copy() + if df is None: + df = pd.DataFrame(columns=error_types[language_type]) + result_stats[task_id]["n_sample"] += 1 + + if "failed" in obj["result"]: + error = obj["result"] + if language_type == "python": + if "assertionerror" in error.lower(): + result_stats[task_id]["assertion error"] += 1 + elif "syntax" in error.lower() or "indent" in error.lower() or "literal" in error.lower(): + result_stats[task_id]["syntax error"] += 1 + elif "not defined" in error.lower(): + result_stats[task_id]["undefined error"] += 1 + elif "timeout" in error.lower(): + result_stats[task_id]["timeout error"] += 1 + elif "type" in error.lower(): + result_stats[task_id]["type error"] += 1 + else: + result_stats[task_id]["runtime error"] += 1 + + elif language_type == "java": + if "wrong answer" in error: + result_stats[task_id]["assertion error"] += 1 + elif "compilation error" in error: + result_stats[task_id]["compilation error"] += 1 + elif "time out" in error: + result_stats[task_id]["timeout error"] += 1 + elif "IndexOutOfBounds" in error: + result_stats[task_id]["index error"] += 1 + elif "UnsupportedOperation" in error: + result_stats[task_id]["unsupported operation"] += 1 + elif "ClassCast" in error: + result_stats[task_id]["class cast error"] += 1 + elif "NullPointer" in error: + result_stats[task_id]["null pointer error"] += 1 + elif "NumberFormat" in error: + result_stats[task_id]["number format error"] += 1 + elif "NoSuchElement" in error: + result_stats[task_id]["no such element"] += 1 + elif "StackOverflow" in error: + result_stats[task_id]["stack overflow"] += 1 + elif "Arithmetic" in error: + result_stats[task_id]["arithmetic error"] += 1 + elif "OutOfMemory" in error: + result_stats[task_id]["out of memory"] += 1 + elif "IllegalArgument" in error: + result_stats[task_id]["illegal argument"] += 1 + else: + result_stats[task_id]["others"] += 1 + + elif language_type == "cpp": + if "compilation error" in error.lower(): + result_stats[task_id]["compilation error"] += 1 + elif "int main(): assertion" in error.lower(): + result_stats[task_id]["assertion error"] += 1 + elif "out_of_range" in error.lower(): + result_stats[task_id]['range error'] += 1 + elif "corrupted top size" in error.lower(): + result_stats[task_id]['range error'] += 1 + elif "length_error" in error.lower(): + result_stats[task_id]['range error'] += 1 + elif "invalid_argument" in error.lower(): + result_stats[task_id]['invalid argument'] += 1 + elif "invalid pointer" in error.lower(): + result_stats[task_id]['pointer error'] += 1 + elif "double free" in error.lower(): + result_stats[task_id]['pointer error'] += 1 + elif "free()" in error.lower(): + result_stats[task_id]['pointer error'] += 1 + elif "logic_error" in error.lower(): + result_stats[task_id]['pointer error'] += 1 + elif "sysmalloc: assertion" in error.lower(): + result_stats[task_id]['pointer error'] += 1 + elif "stack smashing" in error.lower(): + result_stats[task_id]['out of memory'] += 1 + elif "bad_alloc" in error.lower(): + result_stats[task_id]['out of memory'] += 1 + elif "terminate called after throwing an instance of" in error.lower(): + result_stats[task_id]['package error'] += 1 + else: + result_stats[task_id]["others"] += 1 + + elif language_type == "javascript": + if "Assertion failed" in error: + result_stats[task_id]["assertion error"] += 1 + elif "SyntaxError" in error: + result_stats[task_id]["syntax error"] += 1 + elif "ReferenceError" in error: + result_stats[task_id]["undefined error"] += 1 + elif "timed out" in error: + result_stats[task_id]["timeout error"] += 1 + elif "TypeError" in error: + result_stats[task_id]["type error"] += 1 + elif "RangeError" in error: + result_stats[task_id]["range error"] += 1 + else: + result_stats[task_id]["runtime error"] += 1 + + elif language_type == "go": + if "Error: \tNot equal:" in error: + result_stats[task_id]["assertion error"] += 1 + elif "undefined" in error: + result_stats[task_id]["undefined error"] += 1 + elif "expected" in error and "found" in error: + result_stats[task_id]["syntax error"] += 1 + elif "illegal" in error: + result_stats[task_id]["syntax error"] += 1 + elif "unexpected" in error: + result_stats[task_id]["syntax error"] += 1 + elif "FAIL" in error: + result_stats[task_id]["runtime error"] += 1 + elif "timed out" in error: + result_stats[task_id]["timeout error"] += 1 + elif "not used" in error: + result_stats[task_id]['notused error'] += 1 + elif "type" in error: + result_stats[task_id]['type error'] += 1 + else: + incompleted = True + break + else: + if obj["passed"]: + result_stats[task_id]["accepted"] += 1 + + if incompleted: + print(f"Language not supported, aborted. {input_file}") + else: + try: + total, correct = [], [] + for k, res in result_stats.items(): + total.append(res["n_sample"]) + correct.append(res["accepted"]) + df_res = pd.DataFrame(res, index=[int(k.split("/")[-1])]) + df = pd.concat([df, df_res], axis=0) + + total = np.array(total) + correct = np.array(correct) + + ks = [1, 10, 100, 1000] + pass_at_k = {f"pass@{k}": estimate_pass_at_k(total, correct, k).mean() + for k in ks if (total >= k).all()} + + print(pass_at_k) + pass_at_k["file"] = input_file + pass_at_k["n"] = res["n_sample"] + pass_at_k_outs.append(pass_at_k) + + output_prefix = input_file.split("/")[-1].split(".jsonl")[0] + output_file = os.path.join(output_dir, output_prefix + "_stats.xlsx") + df = df.sort_index(ascending=True) + df.to_excel(output_file) + + print(f"Stats saved in {output_file}") + except Exception as e: + print(e) + print(f"Data incompleted, aborted. {input_file}") + + if pass_at_k_outpath is not None: + jsonl_path = os.path.join(output_dir, pass_at_k_outpath) + with open(jsonl_path, "w") as f_out: + for p in pass_at_k_outs: + f_out.write(json.dumps(p) + "\n") + print(f"Pass at k saved in {jsonl_path}") + + +def main(): + fire.Fire(inspect_result) + + +if __name__ == "__main__": + sys.exit(main()) diff --git a/PyTorch/built-in/foundation/CodeGeeX2/benchmark/metric.py b/PyTorch/built-in/foundation/CodeGeeX2/benchmark/metric.py new file mode 100644 index 0000000000000000000000000000000000000000..79d6db3763d0d38e0bb206544fa21ec0d0b6d8c9 --- /dev/null +++ b/PyTorch/built-in/foundation/CodeGeeX2/benchmark/metric.py @@ -0,0 +1,50 @@ +# Copyright (c) OpenAI (https://openai.com) + +# Permission is hereby granted, free of charge, to any person obtaining a copy +# of this software and associated documentation files (the "Software"), to deal +# in the Software without restriction, including without limitation the rights +# to use, copy, modify, merge, publish, distribute, sublicense, and/or sell +# copies of the Software, and to permit persons to whom the Software is +# furnished to do so, subject to the following conditions: + +# The above copyright notice and this permission notice shall be included in +# all copies or substantial portions of the Software. + +# THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR +# IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY, +# FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE +# AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER +# LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM, +# OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN +# THE SOFTWARE. +# ============================================================================ +import itertools +import numpy as np + +from typing import * + + +def estimate_pass_at_k( + num_samples: Union[int, List[int], np.ndarray], + num_correct: Union[List[int], np.ndarray], + k: int +) -> np.ndarray: + """ + Estimates pass@k of each problem and returns them in an array. + """ + + def estimator(n: int, c: int, k: int) -> float: + """ + Calculates 1 - comb(n - c, k) / comb(n, k). + """ + if n - c < k: + return 1.0 + return 1.0 - np.prod(1.0 - k / np.arange(n - c + 1, n + 1)) + + if isinstance(num_samples, int): + num_samples_it = itertools.repeat(num_samples, len(num_correct)) + else: + assert len(num_samples) == len(num_correct) + num_samples_it = iter(num_samples) + + return np.array([estimator(int(n), int(c), k) for n, c in zip(num_samples_it, num_correct)]) diff --git a/PyTorch/built-in/foundation/CodeGeeX2/benchmark/utils.py b/PyTorch/built-in/foundation/CodeGeeX2/benchmark/utils.py new file mode 100644 index 0000000000000000000000000000000000000000..de47797f9806defc817adb8bc10e5f79c5f25fb9 --- /dev/null +++ b/PyTorch/built-in/foundation/CodeGeeX2/benchmark/utils.py @@ -0,0 +1,190 @@ +import os +from typing import * +from codegeex.data.data_utils import stream_jsonl, LANGUAGE_TAG + + +IMPORT_HELPER = { + "python": [ + "import math", + "import re", + "import sys", + "import copy", + "import datetime", + "import itertools", + "import collections", + "import heapq", + "import statistics", + "import functools", + "import hashlib", + "import numpy", + "import numpy as np", + "import string", + "from typing import *", + "from collections import *", + ], + "go" : [ + "math", + "strings", + "fmt", + "strconv", + "time", + "bytes", + "regexp", + "sort", + "math/rand", + "crypto/md5", + ], + "cpp" : [ + "#include", + "#include", + "#include", + "#include", + "#include", + "#include", + "#include", + "#include", + "#include", + ], +} + + +def read_dataset( + data_file: str = None, + dataset_type: str = "humaneval", + num_shot=None, +) -> Dict: + if num_shot is not None: + print(f"{num_shot}-shot setting...") + if "humaneval" in dataset_type.lower(): + if data_file is None: + current_path = os.path.dirname(os.path.abspath(__file__)) + data_file = os.path.join(current_path, "..", "humaneval-x", "python", "data", "humaneval_python.jsonl.gz") + dataset = {task["task_id"]: task for task in stream_jsonl(data_file)} + else: + raise f"Dataset: {dataset_type} not supported." + + return dataset + + +def read_translation_dataset( + data_file_src: str = None, + data_file_tgt: str = None, + lang_src: str = None, + lang_tgt: str = None, + dataset_type: str = "humaneval", +) -> Dict: + if "humaneval" in dataset_type.lower(): + dataset_src = {task["task_id"]: task for task in stream_jsonl(data_file_src)} + dataset_tgt = {task["task_id"].split("/")[-1]: task for task in stream_jsonl(data_file_tgt)} + for k, sample in dataset_src.items(): + prompt = "code translation\n" + if lang_src == "cpp": + prompt += "C++:\n" + elif lang_src == "js": + prompt += "JavaScript:\n" + else: + prompt += f"{lang_src}:\n".capitalize() + prompt += dataset_src[k]["declaration"] + "\n" + dataset_src[k]["canonical_solution"].rstrip() + "\n" + if lang_tgt == "cpp": + prompt += "C++:\n" + elif lang_tgt == "js": + prompt += "JavaScript:\n" + else: + prompt += f"{lang_tgt}:\n".capitalize() + prompt += dataset_tgt[k.split("/")[-1]]["declaration"] + dataset_src[k]["prompt"] = prompt + else: + raise f"Dataset: {dataset_type} not supported." + + return dataset_src + + +def process_extra_prompt(prompt: str, language_type: str = None) -> str: + """ + Processes the extra prompt. + """ + language = language_type.lower() + if language in LANGUAGE_TAG: + extra_prompt = LANGUAGE_TAG[language] + "\n" + else: + extra_prompt = "" + + return extra_prompt + prompt + + +def is_code_generation_finished( + code: str, + language_type: str = None, + dataset: str = None, +): + """ + Checks whether the generated code is finished. + """ + if language_type is None or dataset is None: + return False + + if "humaneval" in dataset.lower(): + if language_type.lower() == "python": + for line in code.split("\n"): + if len(line.strip()) > 0 and line[0] != ' ' and line[0] != '\t': + return True + end_words = ["\ndef", "\nclass", "\nif", "\n#", "\nprint"] + for w in end_words: + if w in code: + return True + elif language_type.lower() == "java": + if code.count("{") + 1 == code.count("}"): + return True + elif language_type.lower() == "go": + if code.count("{") + 1 == code.count("}"): + return True + elif language_type.lower() == "js": + if code.count("{") + 1 == code.count("}"): + return True + elif language_type.lower() == "cpp": + if code.count("{") + 1 == code.count("}"): + return True + + return False + + +def cleanup_code( + code: str, + language_type: str = None, + dataset: str = None, +): + """ + Cleans up the generated code. + """ + if language_type is None or dataset is None: + return code + + if "humaneval" in dataset.lower(): + if language_type.lower() == "python": + end_words = ["\ndef", "\nclass", "\nif", "\n#", "\nprint", "\nassert"] + for w in end_words: + if w in code: + code = code[:code.rfind(w)] + elif language_type.lower() == "java": + main_pos = code.find("public static void main") + if main_pos != -1: + code = code[:main_pos] + '}' + if '}' in code: + code = code[:code.rfind('}')] + '}' + if code.count('{') + 1 == code.count('}'): + code += "\n}" + elif language_type.lower() == "go": + end_words = ["\n//", "\nfunc main("] + for w in end_words: + if w in code: + code = code[:code.rfind(w)] + if '}' in code: + code = code[:code.rfind('}')] + '}' + elif language_type.lower() == "cpp": + if '}' in code: + code = code[:code.rfind('}')] + '}' + elif language_type.lower() == "js": + if '}' in code: + code = code[:code.rfind('}')] + '}' + + return code diff --git a/PyTorch/built-in/foundation/CodeGeeX2/cli_demo.py b/PyTorch/built-in/foundation/CodeGeeX2/cli_demo.py new file mode 100644 index 0000000000000000000000000000000000000000..4aa35d3019313f4dc5d35146b5267457b352b3e9 --- /dev/null +++ b/PyTorch/built-in/foundation/CodeGeeX2/cli_demo.py @@ -0,0 +1,60 @@ +import os +import platform +import signal +from transformers import AutoTokenizer, AutoModel +import readline + +tokenizer = AutoTokenizer.from_pretrained("THUDM/chatglm2-6b", trust_remote_code=True) +model = AutoModel.from_pretrained("THUDM/chatglm2-6b", trust_remote_code=True).cuda() +# 多显卡支持，使用下面两行代替上面一行，将num_gpus改为你实际的显卡数量 +# from utils import load_model_on_gpus +# model = load_model_on_gpus("THUDM/chatglm2-6b", num_gpus=2) +model = model.eval() + +os_name = platform.system() +clear_command = 'cls' if os_name == 'Windows' else 'clear' +stop_stream = False + + +def build_prompt(history): + prompt = "欢迎使用 ChatGLM2-6B 模型，输入内容即可进行对话，clear 清空对话历史，stop 终止程序" + for query, response in history: + prompt += f"\n\n用户：{query}" + prompt += f"\n\nChatGLM2-6B：{response}" + return prompt + + +def signal_handler(signal, frame): + global stop_stream + stop_stream = True + + +def main(): + past_key_values, history = None, [] + global stop_stream + print("欢迎使用 ChatGLM2-6B 模型，输入内容即可进行对话，clear 清空对话历史，stop 终止程序") + while True: + query = input("\n用户：") + if query.strip() == "stop": + break + if query.strip() == "clear": + past_key_values, history = None, [] + os.system(clear_command) + print("欢迎使用 ChatGLM2-6B 模型，输入内容即可进行对话，clear 清空对话历史，stop 终止程序") + continue + print("\nChatGLM：", end="") + current_length = 0 + for response, history, past_key_values in model.stream_chat(tokenizer, query, history=history, + past_key_values=past_key_values, + return_past_key_values=True): + if stop_stream: + stop_stream = False + break + else: + print(response[current_length:], end="", flush=True) + current_length = len(response) + print("") + + +if __name__ == "__main__": + main() diff --git a/PyTorch/built-in/foundation/CodeGeeX2/demo/example_inputs.jsonl b/PyTorch/built-in/foundation/CodeGeeX2/demo/example_inputs.jsonl new file mode 100644 index 0000000000000000000000000000000000000000..848868df208c1479208f7754036f133d3ad58d06 --- /dev/null +++ b/PyTorch/built-in/foundation/CodeGeeX2/demo/example_inputs.jsonl @@ -0,0 +1,9 @@ +{"code": "# Write a quick sort function\n", "langauge": "Python"} +{"code": "// 写一个冒泡排序函数\n", "langauge": "C++"} +{"code": "// 写一个二叉树的类\npublic class", "langauge": "Java"} +{"code": "// 矩阵求行列式\n", "langauge": "Matlab"} +{"code": "\n", "langauge": "HTML"} +{"code": "// 写一个服务器框架， 接收浏览器发过来的请求，并返回处理后的内容\n", "langauge": "JavaScript"} +{"code": "// Write a binary search function\n", "langauge": "Rust"} +{"code": "-- 查询品类最多的三种食品，但是不包含出现时间在2014年及之后\n", "langauge": "SQL"} +{"code": "// Write a simple file system that allows parallel read/write in Golang\n", "langauge": "Go"} \ No newline at end of file diff --git a/PyTorch/built-in/foundation/CodeGeeX2/demo/fastapicpu.py b/PyTorch/built-in/foundation/CodeGeeX2/demo/fastapicpu.py new file mode 100644 index 0000000000000000000000000000000000000000..cd1edab64d917d767ab5ba5c40d92c8b260fc4d3 --- /dev/null +++ b/PyTorch/built-in/foundation/CodeGeeX2/demo/fastapicpu.py @@ -0,0 +1,200 @@ +from fastapi import FastAPI, Request +from transformers import AutoTokenizer, AutoModel +import uvicorn, json, datetime +import torch +import argparse + +try: + import chatglm_cpp + enable_chatglm_cpp = True +except: + print("[WARN] chatglm-cpp not found. Install it by `pip install chatglm-cpp` for better performance. " + "Check out https://github.com/li-plus/chatglm.cpp for more details.") + enable_chatglm_cpp = False + + +#获取选项 +def add_code_generation_args(parser): + group = parser.add_argument_group(title="CodeGeeX2 DEMO") + group.add_argument( + "--model-path", + type=str, + default="THUDM/codegeex2-6b", + ) + group.add_argument( + "--listen", + type=str, + default="127.0.0.1", + ) + group.add_argument( + "--port", + type=int, + default=7860, + ) + group.add_argument( + "--workers", + type=int, + default=1, + ) + group.add_argument( + "--cpu", + action="store_true", + ) + group.add_argument( + "--half", + action="store_true", + ) + group.add_argument( + "--quantize", + type=int, + default=None, + ) + group.add_argument( + "--chatglm-cpp", + action="store_true", + ) + return parser + +LANGUAGE_TAG = { + "Abap" : "* language: Abap", + "ActionScript" : "// language: ActionScript", + "Ada" : "-- language: Ada", + "Agda" : "-- language: Agda", + "ANTLR" : "// language: ANTLR", + "AppleScript" : "-- language: AppleScript", + "Assembly" : "; language: Assembly", + "Augeas" : "// language: Augeas", + "AWK" : "// language: AWK", + "Basic" : "' language: Basic", + "C" : "// language: C", + "C#" : "// language: C#", + "C++" : "// language: C++", + "CMake" : "# language: CMake", + "Cobol" : "// language: Cobol", + "CSS" : "/* language: CSS */", + "CUDA" : "// language: Cuda", + "Dart" : "// language: Dart", + "Delphi" : "{language: Delphi}", + "Dockerfile" : "# language: Dockerfile", + "Elixir" : "# language: Elixir", + "Erlang" : f"% language: Erlang", + "Excel" : "' language: Excel", + "F#" : "// language: F#", + "Fortran" : "!language: Fortran", + "GDScript" : "# language: GDScript", + "GLSL" : "// language: GLSL", + "Go" : "// language: Go", + "Groovy" : "// language: Groovy", + "Haskell" : "-- language: Haskell", + "HTML" : "", + "Isabelle" : "(*language: Isabelle*)", + "Java" : "// language: Java", + "JavaScript" : "// language: JavaScript", + "Julia" : "# language: Julia", + "Kotlin" : "// language: Kotlin", + "Lean" : "-- language: Lean", + "Lisp" : "; language: Lisp", + "Lua" : "// language: Lua", + "Markdown" : "", + "Matlab" : f"% language: Matlab", + "Objective-C" : "// language: Objective-C", + "Objective-C++": "// language: Objective-C++", + "Pascal" : "// language: Pascal", + "Perl" : "# language: Perl", + "PHP" : "// language: PHP", + "PowerShell" : "# language: PowerShell", + "Prolog" : f"% language: Prolog", + "Python" : "# language: Python", + "R" : "# language: R", + "Racket" : "; language: Racket", + "RMarkdown" : "# language: RMarkdown", + "Ruby" : "# language: Ruby", + "Rust" : "// language: Rust", + "Scala" : "// language: Scala", + "Scheme" : "; language: Scheme", + "Shell" : "# language: Shell", + "Solidity" : "// language: Solidity", + "SPARQL" : "# language: SPARQL", + "SQL" : "-- language: SQL", + "Swift" : "// language: swift", + "TeX" : f"% language: TeX", + "Thrift" : "/* language: Thrift */", + "TypeScript" : "// language: TypeScript", + "Vue" : "", + "Verilog" : "// language: Verilog", + "Visual Basic" : "' language: Visual Basic", +} + +app = FastAPI() +def device(): + if enable_chatglm_cpp and args.chatglm_cpp: + print("Using chatglm-cpp to improve performance") + dtype = "f16" if args.half else "f32" + if args.quantize in [4, 5, 8]: + dtype = f"q{args.quantize}_0" + model = chatglm_cpp.Pipeline(args.model_path, dtype=dtype) + return model + + print("chatglm-cpp not enabled, falling back to transformers") + if not args.cpu: + if not args.half: + model = AutoModel.from_pretrained(args.model_path, trust_remote_code=True).cuda() + else: + model = AutoModel.from_pretrained(args.model_path, trust_remote_code=True).cuda().half() + if args.quantize in [4, 8]: + print(f"Model is quantized to INT{args.quantize} format.") + model = model.half().quantize(args.quantize) + else: + model = AutoModel.from_pretrained(args.model_path, trust_remote_code=True) + + return model.eval() + +@app.post("/") +async def create_item(request: Request): + global model, tokenizer + json_post_raw = await request.json() + json_post = json.dumps(json_post_raw) + json_post_list = json.loads(json_post) + lang = json_post_list.get('lang') + prompt = json_post_list.get('prompt') + max_length = json_post_list.get('max_length', 128) + top_p = json_post_list.get('top_p', 0.95) + temperature = json_post_list.get('temperature', 0.2) + top_k = json_post_list.get('top_k', 0) + if lang != "None": + prompt = LANGUAGE_TAG[lang] + "\n" + prompt + if enable_chatglm_cpp and args.chatglm_cpp: + response = model.generate(prompt, + max_length=max_length, + do_sample=temperature > 0, + top_p=top_p, + top_k=top_k, + temperature=temperature) + else: + response = model.chat(tokenizer, + prompt, + max_length=max_length, + top_p=top_p, + top_k=top_k, + temperature=temperature) + now = datetime.datetime.now() + time = now.strftime("%Y-%m-%d %H:%M:%S") + answer = { + "response": response, + "lang": lang, + "status": 200, + "time": time + } + log = "[" + time + "] " + '", prompt:"' + prompt + '", response:"' + repr(response) + '"' + print(log) + + return answer + + +if __name__ == '__main__': + parser = argparse.ArgumentParser() + parser = add_code_generation_args(parser) + args, _ = parser.parse_known_args() + tokenizer = AutoTokenizer.from_pretrained(args.model_path, trust_remote_code=True) + model = device() + uvicorn.run(app, host=args.listen, port=args.port, workers=args.workers) diff --git a/PyTorch/built-in/foundation/CodeGeeX2/demo/gpus.py b/PyTorch/built-in/foundation/CodeGeeX2/demo/gpus.py new file mode 100644 index 0000000000000000000000000000000000000000..025247ac2b8c17ad75d7bf35d28b300bd5e1949e --- /dev/null +++ b/PyTorch/built-in/foundation/CodeGeeX2/demo/gpus.py @@ -0,0 +1,59 @@ +import os +from typing import Dict, Tuple, Union, Optional + +from torch.nn import Module +from transformers import AutoModel + + +def auto_configure_device_map(num_gpus: int) -> Dict[str, int]: + # transformer.word_embeddings 占用1层 + # transformer.final_layernorm 和 lm_head 占用1层 + # transformer.layers 占用 28 层 + # 总共30层分配到num_gpus张卡上 + num_trans_layers = 28 + per_gpu_layers = 30 / num_gpus + + # bugfix: 在linux中调用torch.embedding传入的weight,input不在同一device上,导致RuntimeError + # windows下 model.device 会被设置成 transformer.word_embeddings.device + # linux下 model.device 会被设置成 lm_head.device + # 在调用chat或者stream_chat时,input_ids会被放到model.device上 + # 如果transformer.word_embeddings.device和model.device不同,则会导致RuntimeError + # 因此这里将transformer.word_embeddings,transformer.final_layernorm,lm_head都放到第一张卡上 + # 本文件来源于https://github.com/THUDM/ChatGLM-6B/blob/main/utils.py + # 仅此处做少许修改以支持ChatGLM2,CodeGeeX2 + device_map = { + 'transformer.embedding.word_embeddings': 0, + 'transformer.encoder.final_layernorm': 0, + 'transformer.output_layer': 0, + 'transformer.rotary_pos_emb': 0, + 'lm_head': 0 + } + + used = 2 + gpu_target = 0 + for i in range(num_trans_layers): + if used >= per_gpu_layers: + gpu_target += 1 + used = 0 + assert gpu_target < num_gpus + device_map[f'transformer.encoder.layers.{i}'] = gpu_target + used += 1 + + return device_map + + +def load_model_on_gpus(checkpoint_path: Union[str, os.PathLike], num_gpus: int = 2, + device_map: Optional[Dict[str, int]] = None, **kwargs) -> Module: + if num_gpus < 2 and device_map is None: + model = AutoModel.from_pretrained(checkpoint_path, trust_remote_code=True, **kwargs).half().cuda() + else: + from accelerate import dispatch_model + + model = AutoModel.from_pretrained(checkpoint_path, trust_remote_code=True, **kwargs).half() + + if device_map is None: + device_map = auto_configure_device_map(num_gpus) + + model = dispatch_model(model, device_map=device_map) + + return model diff --git a/PyTorch/built-in/foundation/CodeGeeX2/demo/run_demo.py b/PyTorch/built-in/foundation/CodeGeeX2/demo/run_demo.py new file mode 100644 index 0000000000000000000000000000000000000000..ea58ecb1b1154f7159056f6de4103427066e28ed --- /dev/null +++ b/PyTorch/built-in/foundation/CodeGeeX2/demo/run_demo.py @@ -0,0 +1,364 @@ +import os +import json +import numpy +import torch +import random +import argparse +import gradio as gr + +from transformers import AutoTokenizer, AutoModel + +try: + # Should first install fastllm (https://github.com/ztxz16/fastllm.git) + from fastllm_pytools import llm + enable_fastllm = True +except: + print("fastllm disabled.") + enable_fastllm = False + +try: + from gpus import load_model_on_gpus + enable_multiple_gpus = True +except: + print("Multiple GPUs support disabled.") + enable_multiple_gpus = False + +try: + import chatglm_cpp + enable_chatglm_cpp = True +except: + print("[WARN] chatglm-cpp not found. Install it by `pip install chatglm-cpp` for better performance. " + "Check out https://github.com/li-plus/chatglm.cpp for more details.") + enable_chatglm_cpp = False + + +def get_model(args): + if not args.cpu: + if torch.cuda.is_available(): + device = f"cuda:{args.gpu}" + elif torch.backends.mps.is_built(): + device = "mps" + else: + device = "cpu" + else: + device = "cpu" + + tokenizer = AutoTokenizer.from_pretrained(args.model_path, trust_remote_code=True) + + if args.n_gpus > 1 and enable_multiple_gpus: + # 如需实现多显卡模型加载,传入"n_gpus"为需求的显卡数量 / To enable Multiple GPUs model loading, please adjust "n_gpus" to the desired number of graphics cards. + print(f"Runing on {args.n_gpus} GPUs.") + model = load_model_on_gpus(args.model_path, num_gpus=args.n_gpus) + model = model.eval() + elif enable_chatglm_cpp and args.chatglm_cpp: + print("Using chatglm-cpp to improve performance") + dtype = "f16" + if args.quantize in [4, 5, 8]: + dtype = f"q{args.quantize}_0" + model = chatglm_cpp.Pipeline(args.model_path, dtype=dtype) + else: + model = AutoModel.from_pretrained(args.model_path, trust_remote_code=True) + model = model.eval() + + if enable_fastllm and args.fastllm: + print("fastllm enabled.") + model = model.half() + llm.set_device_map(device) + if args.quantize in [4, 8]: + model = llm.from_hf(model, dtype=f"int{args.quantize}") + else: + model = llm.from_hf(model, dtype="float16") + else: + print("chatglm-cpp and fastllm not installed, using transformers.") + if args.quantize in [4, 8]: + print(f"Model is quantized to INT{args.quantize} format.") + model = model.half().quantize(args.quantize) + model = model.to(device) + + return tokenizer, model + + +def add_code_generation_args(parser): + group = parser.add_argument_group(title="CodeGeeX2 DEMO") + group.add_argument( + "--model-path", + type=str, + default="THUDM/codegeex2-6b", + ) + group.add_argument( + "--example-path", + type=str, + default=None, + ) + group.add_argument( + "--quantize", + type=int, + default=None, + ) + group.add_argument( + "--chatglm-cpp", + action="store_true", + ) + group.add_argument( + "--fastllm", + action="store_true", + ) + group.add_argument( + "--n-gpus", + type=int, + default=1, + ) + group.add_argument( + "--gpu", + type=int, + default=0, + ) + group.add_argument( + "--cpu", + action="store_true", + ) + group.add_argument( + "--listen", + type=str, + default="127.0.0.1", + ) + group.add_argument( + "--port", + type=int, + default=7860, + ) + group.add_argument( + "--username", + type=str, + default=None, + ) + group.add_argument( + "--password", + type=str, + default=None, + ) + group.add_argument( + "--auth", + action="store_true", + ) + + + return parser + + +# 更完编程语言列表请查看 evaluation/utils.py / Full list of supported languages in evaluation/utils.py +LANGUAGE_TAG = { + "Abap" : "* language: Abap", + "ActionScript" : "// language: ActionScript", + "Ada" : "-- language: Ada", + "Agda" : "-- language: Agda", + "ANTLR" : "// language: ANTLR", + "AppleScript" : "-- language: AppleScript", + "Assembly" : "; language: Assembly", + "Augeas" : "// language: Augeas", + "AWK" : "// language: AWK", + "Basic" : "' language: Basic", + "C" : "// language: C", + "C#" : "// language: C#", + "C++" : "// language: C++", + "CMake" : "# language: CMake", + "Cobol" : "// language: Cobol", + "CSS" : "/* language: CSS */", + "CUDA" : "// language: Cuda", + "Dart" : "// language: Dart", + "Delphi" : "{language: Delphi}", + "Dockerfile" : "# language: Dockerfile", + "Elixir" : "# language: Elixir", + "Erlang" : f"% language: Erlang", + "Excel" : "' language: Excel", + "F#" : "// language: F#", + "Fortran" : "!language: Fortran", + "GDScript" : "# language: GDScript", + "GLSL" : "// language: GLSL", + "Go" : "// language: Go", + "Groovy" : "// language: Groovy", + "Haskell" : "-- language: Haskell", + "HTML" : "", + "Isabelle" : "(*language: Isabelle*)", + "Java" : "// language: Java", + "JavaScript" : "// language: JavaScript", + "Julia" : "# language: Julia", + "Kotlin" : "// language: Kotlin", + "Lean" : "-- language: Lean", + "Lisp" : "; language: Lisp", + "Lua" : "// language: Lua", + "Markdown" : "", + "Matlab" : f"% language: Matlab", + "Objective-C" : "// language: Objective-C", + "Objective-C++": "// language: Objective-C++", + "Pascal" : "// language: Pascal", + "Perl" : "# language: Perl", + "PHP" : "// language: PHP", + "PowerShell" : "# language: PowerShell", + "Prolog" : f"% language: Prolog", + "Python" : "# language: Python", + "R" : "# language: R", + "Racket" : "; language: Racket", + "RMarkdown" : "# language: RMarkdown", + "Ruby" : "# language: Ruby", + "Rust" : "// language: Rust", + "Scala" : "// language: Scala", + "Scheme" : "; language: Scheme", + "Shell" : "# language: Shell", + "Solidity" : "// language: Solidity", + "SPARQL" : "# language: SPARQL", + "SQL" : "-- language: SQL", + "Swift" : "// language: swift", + "TeX" : f"% language: TeX", + "Thrift" : "/* language: Thrift */", + "TypeScript" : "// language: TypeScript", + "Vue" : "", + "Verilog" : "// language: Verilog", + "Visual Basic" : "' language: Visual Basic", +} + + +def set_random_seed(seed): + """Set random seed for reproducability.""" + random.seed(seed) + numpy.random.seed(seed) + torch.manual_seed(seed) + + +def main(): + parser = argparse.ArgumentParser() + parser = add_code_generation_args(parser) + args, _ = parser.parse_known_args() + + tokenizer, model = get_model(args) + + examples = [] + if args.example_path is None: + example_path = os.path.join(os.path.split(os.path.realpath(__file__))[0], "example_inputs.jsonl") + else: + example_path = args.example_path + + # Load examples for gradio DEMO + with open(example_path, "r", encoding="utf-8") as f: + for line in f: + examples.append(list(json.loads(line).values())) + + + def predict( + prompt, + lang, + seed, + out_seq_length, + temperature, + top_k, + top_p, + ): + set_random_seed(seed) + if lang != "None": + prompt = LANGUAGE_TAG[lang] + "\n" + prompt + + if enable_fastllm and args.fastllm: + model.direct_query = True + outputs = model.chat(tokenizer, + prompt, + max_length=out_seq_length, + top_p=top_p, + top_k=top_k, + temperature=temperature) + response = prompt + outputs[0] + elif enable_chatglm_cpp and args.chatglm_cpp: + inputs = tokenizer([prompt], return_tensors="pt") + pipeline = model + outputs = pipeline.generate(prompt, + max_length=inputs['input_ids'].shape[-1] + out_seq_length, + do_sample=temperature > 0, + top_p=top_p, + top_k=top_k, + temperature=temperature) + response = prompt + outputs + else: + inputs = tokenizer([prompt], return_tensors="pt") + inputs = inputs.to(model.device) + outputs = model.generate(**inputs, + max_length=inputs['input_ids'].shape[-1] + out_seq_length, + do_sample=True, + top_p=top_p, + top_k=top_k, + temperature=temperature, + pad_token_id=2, + eos_token_id=2) + response = tokenizer.decode(outputs[0]) + + return response + + with gr.Blocks(title="CodeGeeX2 DEMO") as demo: + gr.Markdown( + """ +

+ +

+ """) + gr.Markdown( + """ +

+ 🏠 Homepage｜💻 GitHub｜🛠 Tools VS Code, Jetbrains｜🤗 Download｜📄 Paper +

+ """) + gr.Markdown( + """ + 这是 CodeGeeX2 的简易DEMO。请注意： + * CodeGeeX2 是一个基座模型，它可以完成代码补全/翻译/解释等任务，没有针对聊天进行指令微调。可以在 CodeGeeX 插件[VS Code](https://marketplace.visualstudio.com/items?itemName=aminer.codegeex)、[Jetbrains](https://plugins.jetbrains.com/plugin/20587-codegeex)中体验指令微调后的版本。 + * 可以通过添加`language tag`来控制编程语言，例如`# language: Python`，查看[完整支持语言列表](https://github.com/THUDM/CodeGeeX2/blob/main/evaluation/utils.py#L14)。 + * 按照所选编程语言的格式写注释可以获得更好的结果，请参照下方给出的示例。 + + This is the DEMO for CodeGeeX2. Please note that: + * CodeGeeX2 is a base model, which is not instruction-tuned for chatting. It can do tasks like code completion/translation/explaination. To try the instruction-tuned version in CodeGeeX plugins ([VS Code](https://marketplace.visualstudio.com/items?itemName=aminer.codegeex), [Jetbrains](https://plugins.jetbrains.com/plugin/20587-codegeex)). + * Programming languages can be controled by adding `language tag`, e.g., `# language: Python`. The format should be respected to ensure performance, full list can be found [here](https://github.com/THUDM/CodeGeeX2/blob/main/evaluation/utils.py#L14). + * Write comments under the format of the selected programming language to achieve better results, see examples below. + """) + + with gr.Row(): + with gr.Column(): + prompt = gr.Textbox(lines=14, placeholder='Please enter the description or select an example input below.',label='Input') + with gr.Row(): + gen = gr.Button("Generate") + clr = gr.Button("Clear") + + outputs = gr.Textbox(lines=15, label='Output') + + gr.Markdown( + """ + Generation Parameter + """) + + with gr.Row(): + with gr.Row(): + seed = gr.Slider(maximum=10000, value=8888, step=1, label='Seed') + with gr.Row(): + out_seq_length = gr.Slider(maximum=8192, value=128, minimum=1, step=1, label='Output Sequence Length') + temperature = gr.Slider(maximum=1, value=0.2, minimum=0, label='Temperature') + with gr.Row(): + top_k = gr.Slider(maximum=100, value=0, minimum=0, step=1, label='Top K') + top_p = gr.Slider(maximum=1, value=0.95, minimum=0, label='Top P') + with gr.Row(): + lang = gr.Radio( + choices=["None"] + list(LANGUAGE_TAG.keys()), value='None', label='Programming Language') + inputs = [prompt, lang, seed, out_seq_length, temperature, top_k, top_p] + gen.click(fn=predict, inputs=inputs, outputs=outputs) + clr.click(fn=lambda value: gr.update(value=""), inputs=clr, outputs=prompt) + + gr_examples = gr.Examples(examples=examples, inputs=[prompt, lang], + label="Example Inputs (Click to insert an examplet it into the input box)", + examples_per_page=20) + if not args.auth: + demo.launch(server_name=args.listen, server_port=args.port) + else: + demo.launch(server_name=args.listen, server_port=args.port, auth=(args.username, args.password)) + + #如果需要监听0.0.0.0和其他端口 可以改成 demo.launch(server_name="0.0.0.0", server_port=6666) + #如果需要加密码 demo.launch(server_name="0.0.0.0", server_port=6666, auth=("admin", "password")) + +if __name__ == '__main__': + with torch.no_grad(): + main() + diff --git a/PyTorch/built-in/foundation/CodeGeeX2/docs/zh/inference_zh.md b/PyTorch/built-in/foundation/CodeGeeX2/docs/zh/inference_zh.md new file mode 100644 index 0000000000000000000000000000000000000000..b45580fa8bf5bc74de85ae37c24f847aa9834d50 --- /dev/null +++ b/PyTorch/built-in/foundation/CodeGeeX2/docs/zh/inference_zh.md @@ -0,0 +1,219 @@ +# CodeGeeX2推理教程 + +CodeGeeX2 是多语言代码生成模型 [CodeGeeX](https://github.com/THUDM/CodeGeeX) ([KDD’23](https://arxiv.org/abs/2303.17568)) 的第二代模型，更强，更快，更轻量，是适合本地部署的AI代码生成助手。CodeGeeX2 支持在多种不同平台上进行推理，本教程将会介绍几种不同的推理方式，包括CPU推理，多卡推理，加速推理等。 + +- [快速开始](#快速开始) +- [多精度/量化推理](#多精度/量化推理) +- [多GPU推理](#多GPU推理) +- [Mac推理](#Mac推理) +- [fastllm加速推理](#fastllm加速推理) +- [ChatGLM.cpp量化推理](#chatglmcpp-量化推理) + +## 快速开始 + +下载本仓库并使用`pip`安装环境依赖： + +```shell +git clone https://github.com/THUDM/CodeGeeX2 +cd CodeGeeX2 +pip install -r requirements.txt +``` + +使用`transformers`快速调用[CodeGeeX2-6B](https://huggingface.co/THUDM/codegeex2-6b)，将自动下载权重到本地： + +```python +from transformers import AutoTokenizer, AutoModel +tokenizer = AutoTokenizer.from_pretrained("THUDM/codegeex2-6b", trust_remote_code=True) +model = AutoModel.from_pretrained("THUDM/codegeex2-6b", trust_remote_code=True, device='cuda') # 如使用CPU推理，device='cpu' +model = model.eval() + +# CodeGeeX2支持100种编程语言，加入语言标签引导生成相应的语言 +prompt = "# language: Python\n# write a bubble sort function\n" +inputs = tokenizer.encode(prompt, return_tensors="pt").to(model.device) +outputs = model.generate(inputs, max_length=256, top_k=1) # 示例中使用greedy decoding，检查输出结果是否对齐 +response = tokenizer.decode(outputs[0]) + +>>> print(response) +# language: Python +# write a bubble sort function + + +def bubble_sort(list): + for i in range(len(list) - 1): + for j in range(len(list) - 1): + if list[j] > list[j + 1]: + list[j], list[j + 1] = list[j + 1], list[j] + return list + + +print(bubble_sort([5, 2, 1, 8, 4])) +``` + +亦可以手动下载权重： + +```shell +# huggingface下载 +git clone https://huggingface.co/THUDM/codegeex2-6b +``` + +将tokenizer和model路径改为本地路径： + +```python +model_path = "/path/to/codegeex2-6b" +tokenizer = AutoTokenizer.from_pretrained(model_path, trust_remote_code=True) +model = AutoModel.from_pretrained(model_path, trust_remote_code=True) +``` + +## 多精度/量化推理 + +CodeGeeX2 使用BF16训练，推理时支持BF16/FP16/INT8/INT4，可以根据显卡显存选择合适的精度格式： + +| **Model** | FP16/BF16 | INT8 | INT4 | +| :--------------: | :-------: | :------: | :-----: | +| CodeGeeX-13B | 26\.9 GB | 14\.7 GB | - | +| **CodeGeeX2-6B** | 13\.1 GB | 8\.2 GB | 5\.5 GB | + +默认使用BF16精度进行推理，如显卡不支持BF16（❗️如使用错误的格式，推理结果将出现乱码），需要转换为FP16格式： + +```python +model = AutoModel.from_pretrained(model_path, trust_remote_code=True).half().to("cuda") +``` + +量化推理以INT4为例，可以下载转换好的权重（[INT4权重](https://huggingface.co/THUDM/codegeex2-6b-int4)）或手动转换，如果显卡不支持BF16，也需要先转换为FP16格式： + +```python +# 下载转换好的权重 +model = AutoModel.from_pretrained("THUDM/codegeex2-6b-int4", trust_remote_code=True) + +# 手动转换权重 +model = AutoModel.from_pretrained("THUDM/codegeex2-6b", trust_remote_code=True).quantize(4).to("cuda") + +# 如果显卡不支持BF16，需要先转换为FP16格式 +model = AutoModel.from_pretrained("THUDM/codegeex2-6b", trust_remote_code=True).half().quantize(4).to("cuda") +``` + +## 多GPU推理 + +用[gpus.py](https://github.com/THUDM/CodeGeeX2/blob/main/demo/gpus.py)实现多GPU推理： + +```python +from gpus import load_model_on_gpus +model = load_model_on_gpus("THUDM/codegeex2-6b", num_gpus=2) +``` + +## Mac推理 + +对于搭载了 Apple Silicon 或者 AMD GPU 的 Mac，可以使用 MPS 后端运行。参考 Apple 的 [官方说明](https://developer.apple.com/metal/pytorch) 安装 PyTorch-Nightly（正确的版本号应该是2.x.x.dev2023xxxx，如2.1.0.dev20230729）： + +```shell +pip3 install --pre torch torchvision torchaudio --extra-index-url https://download.pytorch.org/whl/nightly/cpu +``` + +在 MacOS 上只支持从本地加载模型（提前下载权重[codegeex2-6b](https://huggingface.co/THUDM/codegeex2-6b)，[codegeex2-6b-int4](https://huggingface.co/THUDM/codegeex2-6b-int4)），支持FP16/INT8/INT4格式，并使用 mps 后端： + +```python +model = AutoModel.from_pretrained(model_path, trust_remote_code=True).half().to('mps') +``` + +## fastllm加速推理 + +可以使用[fastllm](https://github.com/ztxz16/fastllm)对 CodeGeeX2 进行加速，fastllm是目前支持GLM架构的最快开源框架。首先安装fastllm_pytools： + +```shell +git clone https://github.com/ztxz16/fastllm +cd fastllm +mkdir build +cd build +# 使用GPU编译，需要添加CUDA路径：export CUDA_HOME=/usr/local/cuda/bin:$PATH，export PATH=$PATH:$CUDA_HOME/bin +cmake .. -DUSE_CUDA=ON # 如果不使用GPU编译 cmake .. -DUSE_CUDA=OFF +make -j +cd tools && python setup.py install # 确认安装是否成功，在python中 import fastllm_pytools 不报错 +``` + +如出现架构不支持的报错，需要调整`CMakeLists.txt`，注释掉下面一行： + +```shell +# set(CMAKE_CUDA_ARCHITECTURES "native") +``` +如果是E5系列的CPU可能会出现下面的编译报错 +``` + error: inlining failed in call to ‘always_inline’ ‘__m256i _mm256_add_epi32(__m256i, __m256i)’: target specific option mismatch +``` +此时将'CmakeLists.txt'的第20行修改如下即可编译成功: +``` +set(CMAKE_CXX_FLAGS "${CMAKE_CXX_FLAGS} -pthread --std=c++17 -O2") +``` + +将huggingface转换成fastllm格式： + +```python +# 原本的调用代码 +from transformers import AutoTokenizer, AutoModel +tokenizer = AutoTokenizer.from_pretrained("THUDM/codegeex2-6b", trust_remote_code=True) +model = AutoModel.from_pretrained("THUDM/codegeex2-6b", trust_remote_code=True) + +# 加入下面这两行，将huggingface模型转换成fastllm模型 +from fastllm_pytools import llm +model = llm.from_hf(model, tokenizer, dtype="float16") # dtype支持 "float16", "int8", "int4" +``` + +fastllm中模型接口和huggingface不完全相同，可以参考[demo/run_demo.py](https://github.com/THUDM/CodeGeeX2/blob/main/demo/run_demo.py)中的相关实现： + +```python +model.direct_query = True +outputs = model.chat(tokenizer, + prompt, + max_length=out_seq_length, + top_p=top_p, + top_k=top_k, + temperature=temperature) +response = outputs[0] +``` + +## ChatGLM.cpp 量化推理 + +[ChatGLM.cpp](https://github.com/li-plus/chatglm.cpp) 是类似 LLaMA.cpp 的全平台量化加速方案，支持 q4_0/q4_1/q5_0/q5_1/q8_0 多种量化精度，CPU/CUDA/Metal 多种后端，仅用一行代码实现推理加速。 + +首先安装 chatglm-cpp。如需使用 CUDA 加速，需要添加环境变量 `CMAKE_ARGS="-DGGML_CUBLAS=ON"`；如果仅使用 CPU 加速，将该环境变量去掉即可。 +```sh +CMAKE_ARGS="-DGGML_CUBLAS=ON" pip install chatglm-cpp -v +``` + +仅需一行代码即可量化加速 Hugging Face 模型，`dtype` 可指定 `q4_0`, `q4_1`, `q5_0`, `q5_1`, `q8_0`, `f16`，表示不同的量化类型。 +```python +>>> import chatglm_cpp +>>> +>>> pipeline = chatglm_cpp.Pipeline("THUDM/codegeex2-6b", dtype="q4_0") # Load HF model and quantize it into int4 +Loading checkpoint shards: 100%|███████████████████████████████████████████████| 7/7 [00:09<00:00, 1.33s/it] +Processing model states: 100%|█████████████████████████████████████████████| 199/199 [00:21<00:00, 9.21it/s] +... +>>> print(pipeline.generate("# language: Python\n# write a bubble sort function\n", do_sample=False)) + + +def bubble_sort(list): + for i in range(len(list) - 1): + for j in range(len(list) - 1): + if list[j] > list[j + 1]: + list[j], list[j + 1] = list[j + 1], list[j] + return list + + +print(bubble_sort([5, 4, 3, 2, 1])) +``` + +ChatGLM.cpp 已集成到本仓库，demo 添加选项 `--quantize 4 --chatglm-cpp` 即可开启 int4 (q4_0) 量化加速，例如： +```sh +python ./demo/run_demo.py --quantize 4 --chatglm-cpp +``` + +Fast API 同样支持 ChatGLM.cpp 加速，添加同样参数启动服务： +```sh +python ./demo/fastapicpu.py --quantize 4 --chatglm-cpp +``` + +测试服务接口： +```sh +curl -X POST "http://127.0.0.1:7860" \ + -H 'Content-Type: application/json' \ + -d '{"lang": "Python", "prompt": "# Write a bubble sort function", "max_length": 512}' +``` diff --git a/PyTorch/built-in/foundation/CodeGeeX2/evaluation/README.md b/PyTorch/built-in/foundation/CodeGeeX2/evaluation/README.md new file mode 100644 index 0000000000000000000000000000000000000000..b888e086aa16b4067dddc67eae94bfee8549887e --- /dev/null +++ b/PyTorch/built-in/foundation/CodeGeeX2/evaluation/README.md @@ -0,0 +1,10 @@ +首先从 [Tsinghua Cloud](https://cloud.tsinghua.edu.cn/f/e84444333b6d434ea7b0) 下载处理好的 C-Eval 数据集，解压到 `evaluation` 目录下。然后运行 + +```shell +cd evaluation +python evaluate_ceval.py +``` + +这个脚本会在C-Eval的验证集上进行预测并输出准确率。如果想要得到测试集上的结果可以将代码中的 `./CEval/val/**/*.jsonl` 改为 `./CEval/test/**/*.jsonl`，并按照 C-Eval 规定的格式保存结果并在 [官网](https://cevalbenchmark.com/) 上提交。 + +汇报的结果使用的是内部的并行测试框架，结果可能会有轻微波动。 \ No newline at end of file diff --git a/PyTorch/built-in/foundation/CodeGeeX2/evaluation/__init__.py b/PyTorch/built-in/foundation/CodeGeeX2/evaluation/__init__.py new file mode 100644 index 0000000000000000000000000000000000000000..e69de29bb2d1d6434b8b29ae775ad8c2e48c5391 diff --git a/PyTorch/built-in/foundation/CodeGeeX2/evaluation/evaluate_ceval.py b/PyTorch/built-in/foundation/CodeGeeX2/evaluation/evaluate_ceval.py new file mode 100644 index 0000000000000000000000000000000000000000..bfd317c44d17e6a88383210e0ca6bc1726fc423d --- /dev/null +++ b/PyTorch/built-in/foundation/CodeGeeX2/evaluation/evaluate_ceval.py @@ -0,0 +1,60 @@ +import os +import glob +import re +import json +import torch +import torch.utils.data +from transformers import AutoTokenizer, AutoModel +from tqdm import tqdm + +tokenizer = AutoTokenizer.from_pretrained("THUDM/chatglm2-6b", trust_remote_code=True) +model = AutoModel.from_pretrained("THUDM/chatglm2-6b", trust_remote_code=True).bfloat16().cuda() + +choices = ["A", "B", "C", "D"] +choice_tokens = [tokenizer.encode(choice, add_special_tokens=False)[0] for choice in choices] + + +def build_prompt(text): + return "[Round {}]\n\n问：{}\n\n答：".format(1, text) + + +extraction_prompt = '综上所述，ABCD中正确的选项是：' + +accuracy_dict, count_dict = {}, {} +with torch.no_grad(): + for entry in glob.glob("./CEval/val/**/*.jsonl", recursive=True): + dataset = [] + with open(entry, encoding='utf-8') as file: + for line in file: + dataset.append(json.loads(line)) + correct = 0 + dataloader = torch.utils.data.DataLoader(dataset, batch_size=8) + for batch in tqdm(dataloader): + texts = batch["inputs_pretokenized"] + queries = [build_prompt(query) for query in texts] + inputs = tokenizer(queries, padding=True, return_tensors="pt", truncation=True, max_length=2048).to('cuda') + outputs = model.generate(**inputs, do_sample=False, max_new_tokens=512) + intermediate_outputs = [] + for idx in range(len(outputs)): + output = outputs.tolist()[idx][len(inputs["input_ids"][idx]):] + response = tokenizer.decode(output) + intermediate_outputs.append(response) + answer_texts = [text + intermediate + "\n" + extraction_prompt for text, intermediate in + zip(texts, intermediate_outputs)] + input_tokens = [build_prompt(answer_text) for answer_text in answer_texts] + inputs = tokenizer(input_tokens, padding=True, return_tensors="pt", truncation=True, max_length=2048).to('cuda') + outputs = model(**inputs, return_last_logit=True) + logits = outputs.logits[:, -1] + logits = logits[:, choice_tokens] + preds = logits.argmax(dim=-1) + correct += (preds.cpu() == batch["label"]).sum().item() + accuracy = correct / len(dataset) + print(entry, accuracy) + accuracy_dict[entry] = accuracy + count_dict[entry] = len(dataset) + +acc_total, count_total = 0.0, 0 +for key in accuracy_dict: + acc_total += accuracy_dict[key] * count_dict[key] + count_total += count_dict[key] +print(acc_total / count_total) \ No newline at end of file diff --git a/PyTorch/built-in/foundation/CodeGeeX2/evaluation/evaluation.py b/PyTorch/built-in/foundation/CodeGeeX2/evaluation/evaluation.py new file mode 100644 index 0000000000000000000000000000000000000000..993ec744f27583b0fcad77db232e43eb3847ea9c --- /dev/null +++ b/PyTorch/built-in/foundation/CodeGeeX2/evaluation/evaluation.py @@ -0,0 +1,219 @@ +import os +import re +import sys +import fire +import json +import gzip +import glob +import numpy as np + +from typing import * +from tqdm.auto import tqdm +from collections import defaultdict +from concurrent.futures import ThreadPoolExecutor, as_completed + +from execution import check_correctness +from utils import Logger, IMPORT_HELPER, read_dataset, stream_jsonl_all, estimate_pass_at_k + + +LANGUAGE_NAME = { + "CPP" : "cpp", + "Go" : "go", + "Java" : "java", + "JavaScript" : "js", + "Python" : "python", + "Rust" : "rust", +} + + +def postprocess_generation(sample, generation_mode="completion"): + code = sample["generation"] + if generation_mode == "instruction": + if "```" in code: + pattern = r'```(.*?)\n(.*?)```' + matches = re.findall(pattern, code, re.DOTALL) + for match in matches: + code = match[1] + break + sample["generation"] = code + + return sample + + +def process_test(sample, problems, dataset_type, language_type, generation_mode): + if dataset_type == "humanevalx": + task_id = sample["task_id"] + prompt = problems[task_id]["prompt"] + test = problems[task_id]["test"] + code = sample["generation"] + + # Pre-process for different languages + if language_type == "python": + test_setup = "\n".join(IMPORT_HELPER["python"]) + "\n" + test_string = test_setup + prompt + code + "\n" + test + "\n" + elif language_type == "cpp": + test_set_up = "" + for s in IMPORT_HELPER["cpp"]: + if s not in prompt: + test_set_up += s + "\n" + test_string = test_set_up + "\n" + prompt + code + "\n" + test + elif language_type == "java": + test_string = prompt + code + "\n" + test + elif language_type == "js" or language_type == "javascript": + test_string = prompt + code + "\n" + test + elif language_type == "go": + import_string = problems[task_id]["import"] + prompt = prompt.replace(import_string, "") + test = problems[task_id]["test"] + test_setup = problems[task_id]["test_setup"] + other_pkgs = [] + for pkg in IMPORT_HELPER["go"]: + if pkg not in test_setup: + p = pkg.split("/")[-1] + if p + "." in code: + other_pkgs.append(f"\"{pkg}\"") + if other_pkgs: + import_other_pkgs = "import (\n" + " ".join([p + "\n" for p in other_pkgs]) + ")" + test_string = test_setup + "\n" + import_other_pkgs + "\n" + prompt + code + "\n" + test + else: + test_string = test_setup + "\n" + prompt + code + "\n" + test + elif language_type == "rust": + main = "\nfn main(){ \n } \n" + test_string = main + prompt + code + test + elif dataset_type == "mbpp": + task_id = sample["task_id"] + prompt = sample["prompt"] + test = "\n".join(problems[task_id]["test_list"]) + "\n" + "\n".join(problems[task_id]["challenge_test_list"]) + code = sample["generation"] + test_setup = "\n".join(IMPORT_HELPER["python"]) + "\n" + test_string = test_setup + "\n" + prompt + code + "\n" + problems[task_id]["test_setup_code"] + "\n" + test + "\n" + + return test_string + + +def evaluate_functional_correctness( + input_path: str = None, + output_path: str = None, + log_path: str = None, + tmp_dir: str = "./", + n_workers: int = 32, + timeout: float = 5.0, + k: List[int] = [1, 10, 100], + model_name: str = None, + problem_file: str = None, + language_type: str = None, + dataset_type: str = "humanevalx", + generation_mode: str = "completion", + test_groundtruth: bool = False, +): + if log_path is None: + log_path = os.path.join(output_path, "evaluation.log") + logger = Logger(__name__, log_file=log_path) + + if os.path.isdir(input_path): + input_list = glob.glob(input_path + '/*generation*.jsonl') + sample_jsonl = [] + for input_file in input_list: + sample_jsonl += stream_jsonl_all(input_file) + else: + input_file = input_path + sample_jsonl = stream_jsonl_all(input_file) + + problems = read_dataset(problem_file, dataset_type=dataset_type) + + if output_path is not None: + os.makedirs(output_path, exist_ok=True) + + with ThreadPoolExecutor(max_workers=n_workers) as executor: + + futures = [] + completion_id = Counter() + n_samples = 0 + results = defaultdict(list) + + if test_groundtruth: + logger.info("Testing ground truth...") + else: + logger.info("Testing generation...") + for sample in sample_jsonl: + task_id = sample["task_id"] + if language_type is None: + language_type = LANGUAGE_NAME[task_id.split("/")[0]] + if test_groundtruth: + if dataset_type == "humanevalx": + sample["generation"] = sample["canonical_solution"] + sample["prompt"] = problems[task_id]["prompt"] + if dataset_type == "mbpp": + sample["generation"] = sample["code"] + sample["prompt"] = problems[task_id]["prompt"] + sample = postprocess_generation(sample, generation_mode) + sample["test_code"] = process_test(sample, problems, dataset_type, language_type, generation_mode) + if sample["test_code"] is None: + continue + if "completion_id" in sample: + completion_id_ = sample["completion_id"] + else: + completion_id_ = completion_id[task_id] + args = (task_id, sample, language_type, timeout, tmp_dir, completion_id_) + future = executor.submit(check_correctness, *args) + futures.append(future) + completion_id[task_id] += 1 + n_samples += 1 + + if len(completion_id) == len(problems): + evaluate_pass_at_k = True + else: + evaluate_pass_at_k = False + + logger.info("Running test suites...") + for future in tqdm(as_completed(futures), total=len(futures)): + result = future.result() + results[result["task_id"]].append((result["completion_id"], result)) + + # Calculate pass@k. + total, correct = [], [] + for result in results.values(): + passed = [r[1]["passed"] for r in result] + total.append(len(passed)) + correct.append(sum(passed)) + total = np.array(total) + correct = np.array(correct) + if evaluate_pass_at_k: + ks = k + pass_at_k = {f"pass@{k}": estimate_pass_at_k(total, correct, k).mean() + for k in ks if (total >= k).all()} + logger.info(pass_at_k) + else: + logger.info("Total: {}".format(np.sum(total))) + logger.info("Correct: {}".format(np.sum(correct))) + + if test_groundtruth: + out_file = os.path.join(output_path, "ground_truth.jsonl") + else: + out_file = os.path.join(output_path, "result-" + input_file.split("/")[-2] + "." + input_file.split("/")[-1].split(".")[-1]) + + logger.info("Writing to: {}".format(out_file)) + if out_file.endswith(".gz"): + fp = gzip.GzipFile(fileobj=open(out_file, "wb"), mode="wb") + for res in results.values(): + for r in res: + fp.write((json.dumps(r[1], ensure_ascii=False) + "\n").encode("utf-8")) + else: + fp = open(out_file, 'w') + for res in results.values(): + for r in res: + fp.write(json.dumps(r[1], ensure_ascii=False) + "\n") + fp.close() + + if test_groundtruth: + logger.info("Ground-truth test finished.") + else: + logger.info("Evaluation finished.") + + +def main(): + fire.Fire(evaluate_functional_correctness) + + +if __name__ == "__main__": + sys.exit(main()) diff --git a/PyTorch/built-in/foundation/CodeGeeX2/evaluation/execution.py b/PyTorch/built-in/foundation/CodeGeeX2/evaluation/execution.py new file mode 100644 index 0000000000000000000000000000000000000000..c8b89169614e2c9a58b23efc87d2acc9725de776 --- /dev/null +++ b/PyTorch/built-in/foundation/CodeGeeX2/evaluation/execution.py @@ -0,0 +1,549 @@ +import io +import os +import signal +import random +import gzip +import json +import tempfile +import platform +import subprocess +import contextlib +import faulthandler +import multiprocessing +from typing import * + + +def dicts_to_jsonl(data_list: list, filename: str, compress: bool = True) -> None: + """ + Method saves list of dicts into jsonl file. + :param data: (list) list of dicts to be stored, + :param filename: (str) path to the output file. If suffix .jsonl is not given then methods appends + .jsonl suffix into the file. + :param compress: (bool) should file be compressed into a gzip archive? + """ + sjsonl = '.jsonl' + sgz = '.gz' + # Check filename + if not filename.endswith(sjsonl): + filename = filename + sjsonl + # Save data + + if compress: + filename = filename + sgz + with gzip.open(filename, 'w') as compressed: + for ddict in data_list: + jout = json.dumps(ddict) + '\n' + jout = jout.encode('utf-8') + compressed.write(jout) + else: + with open(filename, 'w') as out: + for ddict in data_list: + jout = json.dumps(ddict) + '\n' + out.write(jout) + + +def check_correctness( + task_id: str, + sample: dict, + language_type: str, + timeout: float = 3.0, + tmp_dir: str = None, + completion_id: Optional[int] = None, +) -> Dict: + """ + Evaluates the functional correctness of a completion by running the test + suite provided in the problem. + """ + + def unsafe_execute(tmp_dir): + random_id = random.uniform(1, 1000) + if "python" in language_type.lower(): + with create_tempdir(): + + # These system calls are needed when cleaning up tempdir. + import os + import shutil + rmtree = shutil.rmtree + rmdir = os.rmdir + chdir = os.chdir + + # Disable functionalities that can make destructive changes to the test. + reliability_guard() + + try: + exec_globals = {} + with swallow_io(): + with time_limit(timeout): + # WARNING + # This program exists to execute untrusted model-generated code. Although + # it is highly unlikely that model-generated code will do something overtly + # malicious in response to this test suite, model-generated code may act + # destructively due to a lack of model capability or alignment. + # Users are strongly encouraged to sandbox this evaluation suite so that it + # does not perform destructive actions on their host or network. + # Once you have read this disclaimer and taken appropriate precautions, + # uncomment the following line and proceed at your own risk: + exec(sample["test_code"], exec_globals) + result.append("passed") + except TimeoutException: + result.append("timed out") + except AssertionError as e: + result.append(f"failed: AssertionError") + except BaseException as e: + result.append(f"failed: {e}") + + # Needed for cleaning up. + shutil.rmtree = rmtree + os.rmdir = rmdir + os.chdir = chdir + + elif "go" in language_type.lower(): + assert tmp_dir is not None, "Go should be evaluated in a dir where necessary module files installed." + + import os + import shutil + + if "tmp" not in tmp_dir: + tmp_dir = os.path.join(tmp_dir, "tmp") + tmp_dir = os.path.join(tmp_dir, f"{task_id.replace('/', '-')}-{random_id}") + if not os.path.exists(tmp_dir): + os.makedirs(tmp_dir) + + os.chdir(tmp_dir) + open(f"main_test.go", 'w').write(sample["test_code"]) + try: + exec_result = None + with time_limit(timeout): + # WARNING + # This program exists to execute untrusted model-generated code. Although + # it is highly unlikely that model-generated code will do something overtly + # malicious in response to this test suite, model-generated code may act + # destructively due to a lack of model capability or alignment. + # Users are strongly encouraged to sandbox this evaluation suite so that it + # does not perform destructive actions on their host or network. + # Once you have read this disclaimer and taken appropriate precautions, + # uncomment the following line and proceed at your own risk: + exec_result = subprocess.run(["go", "test", f"-timeout={timeout}s", "main_test.go"], timeout=timeout, capture_output=True) + + if exec_result.returncode == 0: + result.append("passed") + else: + if exec_result.stderr: + try: + err = exec_result.stderr.decode() + except: + err = exec_result.stderr + else: + try: + err = exec_result.stdout.decode() + except: + err = exec_result.stdout + result.append(f"failed: {err}") + + except TimeoutException: + result.append("timed out") + + shutil.rmtree(tmp_dir) + elif "js" in language_type.lower(): + import os + import shutil + + if "tmp" not in tmp_dir: + tmp_dir = os.path.join(tmp_dir, "tmp") + tmp_dir = os.path.join(tmp_dir, f"{task_id.replace('/', '-')}-{random_id}") + if not os.path.exists(tmp_dir): + os.makedirs(tmp_dir) + + os.chdir(tmp_dir) + open(f"test.js", 'w').write(sample["test_code"]) + try: + exec_result = None + with time_limit(timeout): + # WARNING + # This program exists to execute untrusted model-generated code. Although + # it is highly unlikely that model-generated code will do something overtly + # malicious in response to this test suite, model-generated code may act + # destructively due to a lack of model capability or alignment. + # Users are strongly encouraged to sandbox this evaluation suite so that it + # does not perform destructive actions on their host or network. + # Once you have read this disclaimer and taken appropriate precautions, + # uncomment the following line and proceed at your own risk: + exec_result = subprocess.run(["node", "test.js"], timeout=timeout, capture_output=True) + + if exec_result.stderr.decode(): + err = exec_result.stderr.decode() + result.append(f"failed: {err}") + elif exec_result.stdout.decode(): + err = exec_result.stdout.decode() + result.append(f"failed: {err}") + else: + result.append("passed") + + except TimeoutException: + result.append("timed out") + + shutil.rmtree(tmp_dir) + elif "cpp" in language_type.lower(): + import os + import shutil + + if "tmp" not in tmp_dir: + tmp_dir = os.path.join(tmp_dir, "tmp") + tmp_dir = os.path.join(tmp_dir, f"{task_id.replace('/', '-')}-{random_id}") + if not os.path.exists(tmp_dir): + os.makedirs(tmp_dir) + + os.chdir(tmp_dir) + open(f"test.cpp", 'w').write(sample["test_code"]) + if "162" in task_id: + compilation_result = subprocess.run(["/usr/bin/g++", "-std=c++11", "test.cpp", "-lcrypto", "-lssl"], + timeout=timeout, + capture_output=True) + else: + compilation_result = subprocess.run(["/usr/bin/g++", "-std=c++11", "test.cpp"], timeout=timeout, + capture_output=True) + if compilation_result.returncode != 0: + if compilation_result.stderr: + err = compilation_result.stderr.decode() + else: + err = compilation_result.stdout.decode() + result.append(f"failed: compilation error: {err}") + else: + try: + exec_result = None + with time_limit(timeout): + # WARNING + # This program exists to execute untrusted model-generated code. Although + # it is highly unlikely that model-generated code will do something overtly + # malicious in response to this test suite, model-generated code may act + # destructively due to a lack of model capability or alignment. + # Users are strongly encouraged to sandbox this evaluation suite so that it + # does not perform destructive actions on their host or network. + # Once you have read this disclaimer and taken appropriate precautions, + # uncomment the following line and proceed at your own risk: + exec_result = subprocess.run(["./a.out"], timeout=timeout, capture_output=True) + + if exec_result.returncode == 0: + result.append("passed") + else: + if exec_result.stderr: + try: + err = exec_result.stderr.decode() + except: + err = exec_result.stderr + else: + try: + err = exec_result.stdout.decode() + except: + err = exec_result.stdout + result.append(f"failed: {err}") + except TimeoutException: + result.append("timed out") + + shutil.rmtree(tmp_dir) + elif "rust" in language_type.lower(): + import os + WD: str = os.path.dirname(tmp_dir) + RUST_DIR: str = os.path.join(WD, "rust") + RUST_SRC: str = os.path.join(RUST_DIR, "src") + RUST_BIN: str = os.path.join(RUST_SRC, "bin") + RUST_TMP_DIR: str = os.path.join(RUST_DIR, "tmp") + RUST_LOGS: str = os.path.join(RUST_TMP_DIR, "logs") + RUST_EXT: str = ".rs" + + # Create mandatory tmp directories + os.makedirs(RUST_TMP_DIR, exist_ok=True) + os.makedirs(RUST_LOGS, exist_ok=True) + os.makedirs(RUST_SRC, exist_ok=True) + os.makedirs(RUST_BIN, exist_ok=True) + + with tempfile.NamedTemporaryFile(dir = RUST_BIN, delete=False) as f: + # temporal file name + file_prefix = sample["task_id"].lower().replace("/", "_") + file_name:str = file_prefix +RUST_EXT + + os.rename(f.name, os.path.join(RUST_BIN, file_name)) + + # Sample to pure Rust function + rust_code: str = sample["test_code"] + + # dump the rust source code in the target temporal file + f.write(rust_code.encode('utf-8')) + + # Proceed towards Rust binaries compilation. Therefore move to Rust module root dir. + os.chdir(RUST_DIR) + + # Two possible outcomes + # Pass OR Fail compilation + log_filename: str = file_prefix + ".jsonl" + log_path: str = os.path.join(RUST_LOGS, log_filename) + cargo_check: str = "cargo check --bin " + file_prefix + " --message-format json >> " + log_path + # Compilation build status + returned_val_compilation: int + + # Overwrite file content + if os.path.exists(log_path): + if(file_size := os.path.getsize(log_path)) >= 0: + os.remove(log_path) + returned_val_compilation = os.system(cargo_check) + + else: + returned_val_compilation = os.system(cargo_check) + + # 0 means success + if returned_val_compilation == 0: + + #Execution pipeline + cargo_test: str = "cargo test --bin " +file_prefix+ " --message-format json >> " + log_path + returned_val_execution = os.system(cargo_test) + + if returned_val_execution == 0: + result.append("passed") + else: + result.append(f"failed: execution error") + + else: + result.append(f"failed: compilation error") + + + elif "java" in language_type.lower(): + assert tmp_dir is not None, "Java should be evaluated in a temporary dir." + + import os + import shutil + + if "tmp" not in tmp_dir: + tmp_dir = os.path.join(tmp_dir, "tmp") + tmp_dir = os.path.join(tmp_dir, f"{task_id.replace('/', '-')}-{random_id}") + if not os.path.exists(tmp_dir): + os.makedirs(tmp_dir) + + os.chdir(tmp_dir) + open(os.path.join(tmp_dir, "Main.java"), 'w').write(sample["test_code"]) + res = "failed: unknown error" + compile_returncode = -1 + for _ in range(5): + try: + compilation_result = subprocess.run(['javac', os.path.join(tmp_dir, "Main.java")], timeout=5, + capture_output=True) + compile_returncode = compilation_result.returncode + break + except subprocess.TimeoutExpired as e: + continue + if compile_returncode != 0: + res = "failed: compilation error" + else: + exec_result = None + try: + # WARNING + # This program exists to execute untrusted model-generated code. Although + # it is highly unlikely that model-generated code will do something overtly + # malicious in response to this test suite, model-generated code may act + # destructively due to a lack of model capability or alignment. + # Users are strongly encouraged to sandbox this evaluation suite so that it + # does not perform destructive actions on their host or network. + # Once you have read this disclaimer and taken appropriate precautions, + # uncomment the following line and proceed at your own risk: + exec_result = subprocess.run([f'java', '-cp', tmp_dir, 'Main'], timeout=timeout, capture_output=True) + if exec_result.returncode == 0: + res = "passed" + elif exec_result.returncode == 1: + if "AssertionError" in exec_result.stderr.decode('unicode-escape'): + res = "failed: wrong answer" + else: + res = f"failed: {exec_result.stderr.decode()}" + except subprocess.TimeoutExpired as e: + res = "time out" + except BaseException as e: + res = f"failed: {e}" + result.append(res) + + shutil.rmtree(tmp_dir) + + manager = multiprocessing.Manager() + result = manager.list() + + p = multiprocessing.Process(target=unsafe_execute, args=(tmp_dir,)) + p.start() + p.join(timeout=timeout + 1) + if p.is_alive(): + p.kill() + + if not result: + result.append("timed out") + + return { + "task_id" : task_id, + "completion_id": completion_id, + "test_code" : sample["test_code"], + "prompt" : sample["prompt"], + "generation" : sample["generation"], + "result" : result[0], + "passed" : result[0] == "passed", + "finish" : -1 if "finish" not in sample else sample["finish"], + "file" : "" if "file" not in sample else sample["file"], + "output" : [] if "output" not in sample else sample["output"], + } + +# Copyright (c) OpenAI (https://openai.com) + +# Permission is hereby granted, free of charge, to any person obtaining a copy +# of this software and associated documentation files (the "Software"), to deal +# in the Software without restriction, including without limitation the rights +# to use, copy, modify, merge, publish, distribute, sublicense, and/or sell +# copies of the Software, and to permit persons to whom the Software is +# furnished to do so, subject to the following conditions: + +# The above copyright notice and this permission notice shall be included in +# all copies or substantial portions of the Software. + +# THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR +# IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY, +# FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE +# AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER +# LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM, +# OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN +# THE SOFTWARE. +# ============================================================================ +@contextlib.contextmanager +def time_limit(seconds: float): + def signal_handler(signum, frame): + raise TimeoutException("Timed out!") + + signal.setitimer(signal.ITIMER_REAL, seconds) + signal.signal(signal.SIGALRM, signal_handler) + try: + yield + finally: + signal.setitimer(signal.ITIMER_REAL, 0) + + +@contextlib.contextmanager +def swallow_io(): + stream = WriteOnlyStringIO() + with contextlib.redirect_stdout(stream): + with contextlib.redirect_stderr(stream): + with redirect_stdin(stream): + yield + + +@contextlib.contextmanager +def create_tempdir(): + with tempfile.TemporaryDirectory() as dirname: + with chdir(dirname): + yield dirname + + +class TimeoutException(Exception): + pass + + +class WriteOnlyStringIO(io.StringIO): + """ StringIO that throws an exception when it's read from """ + + def read(self, *args, **kwargs): + raise IOError + + def readline(self, *args, **kwargs): + raise IOError + + def readlines(self, *args, **kwargs): + raise IOError + + def readable(self, *args, **kwargs): + """ Returns True if the IO object can be read. """ + return False + + +class redirect_stdin(contextlib._RedirectStream): # type: ignore + _stream = 'stdin' + + +@contextlib.contextmanager +def chdir(root): + if root == ".": + yield + return + cwd = os.getcwd() + os.chdir(root) + try: + yield + except BaseException as exc: + raise exc + finally: + os.chdir(cwd) + + +def reliability_guard(maximum_memory_bytes: Optional[int] = None): + """ + This disables various destructive functions and prevents the generated code + from interfering with the test (e.g. fork bomb, killing other processes, + removing filesystem files, etc.) + + WARNING + This function is NOT a security sandbox. Untrusted code, including, model- + generated code, should not be blindly executed outside of one. See the + Codex paper for more information about OpenAI's code sandbox, and proceed + with caution. + """ + + if maximum_memory_bytes is not None: + import resource + resource.setrlimit(resource.RLIMIT_AS, (maximum_memory_bytes, maximum_memory_bytes)) + resource.setrlimit(resource.RLIMIT_DATA, (maximum_memory_bytes, maximum_memory_bytes)) + if not platform.uname().system == 'Darwin': + resource.setrlimit(resource.RLIMIT_STACK, (maximum_memory_bytes, maximum_memory_bytes)) + + faulthandler.disable() + + import builtins + builtins.exit = None + builtins.quit = None + + import os + os.environ['OMP_NUM_THREADS'] = '1' + + os.kill = None + os.system = None + os.putenv = None + os.remove = None + os.removedirs = None + os.rmdir = None + os.fchdir = None + os.setuid = None + os.fork = None + os.forkpty = None + os.killpg = None + os.rename = None + os.renames = None + os.truncate = None + os.replace = None + os.unlink = None + os.fchmod = None + os.fchown = None + os.chmod = None + os.chown = None + os.chroot = None + os.fchdir = None + os.lchflags = None + os.lchmod = None + os.lchown = None + os.getcwd = None + os.chdir = None + + import shutil + shutil.rmtree = None + shutil.move = None + shutil.chown = None + + import subprocess + subprocess.Popen = None # type: ignore + + __builtins__['help'] = None + + import sys + sys.modules['ipdb'] = None + sys.modules['joblib'] = None + sys.modules['resource'] = None + sys.modules['psutil'] = None + sys.modules['tkinter'] = None diff --git a/PyTorch/built-in/foundation/CodeGeeX2/evaluation/generation.py b/PyTorch/built-in/foundation/CodeGeeX2/evaluation/generation.py new file mode 100644 index 0000000000000000000000000000000000000000..3b1bca72261d6e2b7b0961c6f07e2f4266022766 --- /dev/null +++ b/PyTorch/built-in/foundation/CodeGeeX2/evaluation/generation.py @@ -0,0 +1,439 @@ +import os +import zmq +import time +import json +import torch +import random +import socket +import argparse + +from typing import * +from transformers import AutoModel, AutoModelForCausalLM, AutoTokenizer, StoppingCriteria +from utils import Logger, read_dataset, process_extra_prompt, is_code_generation_finished, cleanup_code + +logger = Logger(__name__) + + +def add_code_generation_specific_args(parser): + group = parser.add_argument_group("Code Generation") + group.add_argument( + "--hostfile", + type=str, + default="./hostfile", + ) + group.add_argument( + "--channel-ip", + type=str, + default=None, + help="IP for ZeroMQ channel", + ) + group.add_argument( + "--channel-port", + type=int, + default=5555, + help="Port for ZeroMQ channel", + ) + group.add_argument( + "--master-port", + type=int, + default=6007, + help="Port for distributed channel", + ) + group.add_argument( + "--model-per-device", + type=int, + default=1, + help="Number of models per device", + ) + group.add_argument( + "--max-length", + type=int, + default=8192, + help="Max sequence length", + ) + group.add_argument( + "--top-p", + type=float, + default=1.0, + help="Top-p Probability for sampling", + ) + group.add_argument( + "--top-k", + type=int, + default=0, + help="Top-k for sampling", + ) + group.add_argument( + "--temperature", + type=float, + default=1.0, + help="Temperature for sampling", + ) + group.add_argument( + "--greedy", + type=int, + default=0, + help="Use greedy decoding instead of sampling", + ) + group.add_argument( + "--seed", + type=int, + default=42, + help="Random seed", + ) + group.add_argument( + "--micro-batch-size", + type=int, + default=1, + help="Micro batch size for each GPU", + ) + group.add_argument( + "--samples-per-problem", + type=int, + default=200, + help="Number of samples to generate for each problem", + ) + group.add_argument( + "--gen-node-world-size", + type=int, + default=1, + help="Number of machines to use for generation", + ) + group.add_argument( + '--task-name', + default="generation", + help='Name of task', + ) + group.add_argument( + '--model-name', + default="codegeex2-6b", + help='Name of model, support ["codegeex2-6b", "starcoder", "replit-code-v1-3b", "codegen25-7b-multi", "codegen25-7b-mono", "codegen-16B-multi"]', + ) + group.add_argument( + '--data-path', + required=True, + ) + group.add_argument( + '--output-path', + required=True, + ) + group.add_argument( + '--log-path', + default=None, + help='Path to log output', + ) + group.add_argument( + '--model-path', + required=True, + ) + group.add_argument( + '--dataset-type', + default="humanevalx", + help='Identify the evaluation dataset [humanevalx]', + ) + group.add_argument( + '--language-type', + default="python", + help='Identify the type of programming language to generate', + ) + group.add_argument( + '--generation-mode', + default="instruction", + ) + + +class CodeStoppingCriteria(StoppingCriteria): + """ + This class can be used to stop generation whenever the full generated number of tokens exceeds `max_length` or meet the code generation stopping criteria. + """ + + def __init__( + self, + max_length: int, + micro_batch_size: int, + tokenizer, + dataset_type: str, + language_type: str, + prompt: str, + ): + self.max_length = max_length + self.tokenizer = tokenizer + self.dataset_type = dataset_type + self.language_type = language_type + self.prompt = prompt + self.stop_index = [-1 for _ in range(micro_batch_size)] + + def __call__(self, input_ids: torch.LongTensor, scores: torch.FloatTensor, **kwargs) -> bool: + for i, input_id in enumerate(input_ids): + if self.stop_index[i] > -1: + continue + code = self.tokenizer.decode(input_id) + code = code[len(self.prompt):] + if is_code_generation_finished( + code, + dataset_type=self.dataset_type, + language_type=self.language_type) or input_id.shape[-1] >= self.max_length: + self.stop_index[i] = len(code) + len(self.prompt) + if all([s != -1 for s in self.stop_index]): + return True + + return False + + +def run_generation_distributed(args, model, tokenizer): + logger.info(f"Connecting to tcp://{args.channel_ip}:{args.channel_port}") + context = zmq.Context() + socket = context.socket(zmq.REQ) + socket.connect(f"tcp://{args.channel_ip}:{args.channel_port}") + + os.makedirs(args.output_path, exist_ok=True) + output_path = os.path.join( + args.output_path, + f"{args.task_name}-t{args.temperature}-topp{args.top_p}-ns{args.samples_per_problem}-rank{args.rank}.jsonl", + ) + + def process(obj): + results = [] + prompt = obj["prompt"] + if args.generation_mode == "instruction": + inputs = tokenizer([prompt] * args.micro_batch_size, return_tensors="pt") + inputs = inputs.to(model.device) + outputs = model.generate(**inputs, + max_length=args.max_length, + do_sample=True if not args.greedy else False, + use_cache=True, + top_p=args.top_p, + top_k=args.top_k, + temperature=args.temperature, + pad_token_id=tokenizer.eos_token_id) + for i, output in enumerate(outputs): + response = tokenizer.decode(output) + res = obj.copy() + res["generation"] = response[len(prompt):].strip() + results.append(res) + elif args.generation_mode == "completion": + inputs = tokenizer([prompt for _ in range(args.micro_batch_size)], return_tensors="pt") + inputs = inputs.to(model.device) + stop_criteria = CodeStoppingCriteria( + max_length=args.max_length, + micro_batch_size=args.micro_batch_size, + tokenizer=tokenizer, + dataset_type=args.dataset_type, + language_type=args.language_type, + prompt=prompt) + outputs = model.generate(**inputs, + max_length=args.max_length, + do_sample=True if not args.greedy else False, + use_cache=True, + stopping_criteria=[stop_criteria], + top_p=args.top_p, + top_k=args.top_k, + temperature=args.temperature, + pad_token_id=tokenizer.eos_token_id) + for i, output in enumerate(outputs): + response = tokenizer.decode(output) + res = obj.copy() + res["generation_raw"] = response + res["generation"] = cleanup_code( + response[len(prompt):], + dataset_type=args.dataset_type, + language_type=args.language_type) + results.append(res) + + return results + + fout = open(output_path, "w", encoding="utf-8") + while True: + socket.send_json({"rank": args.rank, "action": "pull"}) + resp = socket.recv_json() + try: + if resp["task_id"] is None: + break + + current_spec = resp["task_id"] + results = process(current_spec) + + for res in results: + fout.write(json.dumps(res, ensure_ascii=False) + "\n") + fout.flush() + + socket.send_json( + { + "rank" : args.rank, + "action" : "success", + "task_id": current_spec['task_id'] + } + ) + socket.recv() + + except Exception as e: + logger.error(f"*** (rank={args.rank}) crashed.") + logger.error(f" error: {repr(e)}") + socket.send_json( + { + "rank" : args.rank, + "action" : "fail", + "task_id": current_spec['task_id'] + } + ) + socket.recv() + continue + + +def main(args, node_rank: int, local_rank: int, master_port: int, num_devices: int): + world_size = args.gen_node_world_size * num_devices + args.rank = num_devices * node_rank + local_rank + args.world_size = world_size + logger.info(f"Generating on rank {args.rank} of {args.world_size}") + + try: + if args.model_name in ["codegeex2-6b"]: + tokenizer = AutoTokenizer.from_pretrained(args.model_path, trust_remote_code=True) + else: + tokenizer = AutoTokenizer.from_pretrained(args.model_path, clean_up_tokenization_spaces=False, trust_remote_code=True) + if args.model_name in ["codegeex2-6b"]: + model = AutoModel.from_pretrained(args.model_path, trust_remote_code=True).to("cuda:{}".format(local_rank % torch.cuda.device_count())) + elif args.model_name in ["starcoder", "replit-code-v1-3b", "codegen25-7b-multi", "codegen25-7b-mono", "codegen-16B-multi"]: + model = AutoModelForCausalLM.from_pretrained(args.model_path, trust_remote_code=True).to("cuda:{}".format(local_rank % torch.cuda.device_count())) + else: + try: + model = AutoModel.from_pretrained(args.model_path, trust_remote_code=True).to("cuda:{}".format(local_rank % torch.cuda.device_count())) + except: + logger.error(f"Model {args.model_name} not supported.") + raise NotImplementedError + except Exception as e: + logger.error(e) + + model = model.eval() + # Generate samples. + run_generation_distributed(args, model, tokenizer) + + logger.info(f"rank={args.rank} worker finished, waiting ...") + exit(0) + + +def server(args): + logger.info(f"[ server ] starting ...") + entries = read_dataset(args.data_path, dataset_type=args.dataset_type) + + assert args.samples_per_problem % args.micro_batch_size == 0, "samples_per_problem should be divisible by batch_size" + + for entry in entries.values(): + entry["prompt"] = process_extra_prompt( + entry["prompt"], + language_type=args.language_type, + dataset_type=args.dataset_type, + generation_mode=args.generation_mode, + ) + + res = [] + for entry in entries.values(): + res.extend([entry] * (args.samples_per_problem // args.micro_batch_size)) + random.shuffle(res) + all_entries = res + + # setup zeromq channel + logger.info(f"[ server ] starting up on port {args.channel_port}") + context = zmq.Context() + logger.info(f"[ server ] creating socket") + socket = context.socket(zmq.REP) + logger.info(f"[ server ] binding to port {args.channel_port}") + socket.bind(f"tcp://*:{args.channel_port}") + + logger.info( + f"[ server ] loaded {len(entries)} entries, generating {len(entries) * args.samples_per_problem} samples", + ) + + remaining_entries = all_entries.copy() + running_workers = args.gen_node_world_size * torch.cuda.device_count() + num_finished = 0 + + logger.info(f"[ server ] listening for requests ...") + start_time = time.perf_counter() + while True: + # Wait for next request from client + msg = socket.recv_json() + rank = msg["rank"] + action = msg["action"] + + if action == "pull": + if len(remaining_entries) == 0: + socket.send_json({"task_id": None}) + running_workers -= 1 + logger.info(f"[ server ] Shutting down worker {rank}, remaining {running_workers} workers") + if running_workers == 0 and num_finished == len(all_entries): + logger.info(f"[ server ] All workers finished") + break + else: + entry = remaining_entries.pop() + time_elapsed = time.perf_counter() - start_time + logger.info(f"[ server ] Sending entry {entry['task_id']} to worker {rank}") + remaining = ( + len(remaining_entries) + / (len(all_entries) - len(remaining_entries)) + * time_elapsed + ) + time_per_sampple = 0.0 if num_finished == 0 else time_elapsed / num_finished / args.micro_batch_size + logger.info( + f"[ server ] total {len(all_entries)}, assigned {len(all_entries) - len(remaining_entries)}, finished {num_finished}, elapsed {time_elapsed:.4f}, speed {time_per_sampple:.4f}s/sample, remaining {remaining:.4f}", + ) + socket.send_json({"task_id": entry}) + else: + if action == "success": + logger.info(f"[ server ] {msg['task_id']} is finished") + socket.send_json({"pong": 1}) + else: + logger.info(f"[ server ] {msg['task_id']} is not finished") + remaining_entries.append(msg['task_id']) + socket.send_json({"pong": 1}) + break + + num_finished += 1 + + +if __name__ == "__main__": + torch.multiprocessing.set_start_method("spawn") + parser = argparse.ArgumentParser() + add_code_generation_specific_args(parser) + args = parser.parse_args() + + if args.log_path is None: + args.log_path = os.path.join(args.output_path, "generation.log") + + logger.info("start method: " + torch.multiprocessing.get_start_method()) + + processes = [] + num_devices = torch.cuda.device_count() + hosts = open(args.hostfile, "r").readlines() + hosts = [host.strip() for host in hosts] + master_port = args.master_port + + node_rank = None + for i in range(len(hosts)): + if hosts[i] == socket.gethostbyname(socket.gethostname()): + node_rank = i + break + assert ( + node_rank is not None + ), f"Could not find hostname ({socket.gethostbyname(socket.gethostname())}) in hostlist" + + # launch server + if socket.gethostbyname(socket.gethostname()) == hosts[0]: + server_process = torch.multiprocessing.Process(target=server, args=(args,)) + logger.info(f"Launching server ...") + server_process.start() + processes.append(server_process) + + for i in range(num_devices): + local_rank = i + logger.info(f"launching local rank {i}") + + p = torch.multiprocessing.Process( + target=main, + args=(args, node_rank, local_rank, master_port, num_devices), + ) + p.start() + processes.append(p) + + for p in processes: + p.join() diff --git a/PyTorch/built-in/foundation/CodeGeeX2/evaluation/inspect_jsonl.py b/PyTorch/built-in/foundation/CodeGeeX2/evaluation/inspect_jsonl.py new file mode 100644 index 0000000000000000000000000000000000000000..9ee73250132c68750c6444cdcc210e749c99aef1 --- /dev/null +++ b/PyTorch/built-in/foundation/CodeGeeX2/evaluation/inspect_jsonl.py @@ -0,0 +1,46 @@ +import fire +import json +import numpy as np + +from typing import * +from utils import Logger + + +def main( + data_path: str = "./test.jsonl", + threshold: int = -1, + random: int = 0, + log_path: str = 'inspect_jsonl.txt', + random_rate: float = 0.5, +): + logger = Logger(__name__, log_file=log_path, log_mode="file", disable_formatter=True) + + n = 0 + with open(data_path, "r") as f: + for i, line in enumerate(f): + if i == 0: + logger.info("Data has the following keys") + obj = json.loads(line) + logger.info(obj.keys()) + if threshold > 0 and n > threshold: + break + if random and np.random.randint(10) > 10 * random_rate: + continue + + obj = json.loads(line) + n += 1 + logger.info(f"========== Sample {i} ==========") + if 'code' in obj: + try: + code_splits = obj['code'].split("\n") + logger.info(f"Length of chars: {len(obj['code'])}, length of lines: {len(code_splits)}.") + except: + pass + for j, k in enumerate(obj.keys()): + logger.info(f"** Key {j}: {k} **") + logger.info(obj[k]) + print(f"Log saved in {log_path}") + + +if __name__ == "__main__": + fire.Fire(main) \ No newline at end of file diff --git a/PyTorch/built-in/foundation/CodeGeeX2/evaluation/utils.py b/PyTorch/built-in/foundation/CodeGeeX2/evaluation/utils.py new file mode 100644 index 0000000000000000000000000000000000000000..0147bc0ec78042c29395cc6f2a7cfea92bceee42 --- /dev/null +++ b/PyTorch/built-in/foundation/CodeGeeX2/evaluation/utils.py @@ -0,0 +1,501 @@ +import re +import json +import gzip +import torch +import numpy +import random +import logging +import itertools +import numpy as np + +from typing import * + + +LANGUAGE_TAG = { + "c" : "// language: C", + "c++" : "// language: C++", + "cpp" : "// language: C++", + "c#" : "// language: C#", + "csharp" : "// language: C#", + "c-sharp" : "// language: C#", + "css" : "/* language: CSS */", + "cuda" : "// language: Cuda", + "dart" : "// language: Dart", + "lua" : "// language: Lua", + "objectivec" : "// language: Objective-C", + "objective-c" : "// language: Objective-C", + "objective-c++": "// language: Objective-C++", + "python" : "# language: Python", + "perl" : "# language: Perl", + "prolog" : f"% language: Prolog", + "swift" : "// language: swift", + "lisp" : "; language: Lisp", + "java" : "// language: Java", + "scala" : "// language: Scala", + "tex" : f"% language: TeX", + "vue" : "", + "markdown" : "", + "html" : "", + "php" : "// language: PHP", + "js" : "// language: JavaScript", + "javascript" : "// language: JavaScript", + "typescript" : "// language: TypeScript", + "go" : "// language: Go", + "shell" : "# language: Shell", + "rust" : "// language: Rust", + "sql" : "-- language: SQL", + "kotlin" : "// language: Kotlin", + "vb" : "' language: Visual Basic", + "ruby" : "# language: Ruby", + "pascal" : "// language: Pascal", + "r" : "# language: R", + "fortran" : "!language: Fortran", + "lean" : "-- language: Lean", + "matlab" : f"% language: Matlab", + "delphi" : "{language: Delphi}", + "scheme" : "; language: Scheme", + "basic" : "' language: Basic", + "assembly" : "; language: Assembly", + "groovy" : "// language: Groovy", + "abap" : "* language: Abap", + "gdscript" : "# language: GDScript", + "haskell" : "-- language: Haskell", + "julia" : "# language: Julia", + "elixir" : "# language: Elixir", + "excel" : "' language: Excel", + "clojure" : "; language: Clojure", + "actionscript" : "// language: ActionScript", + "solidity" : "// language: Solidity", + "powershell" : "# language: PowerShell", + "erlang" : f"% language: Erlang", + "cobol" : "// language: Cobol", + "alloy" : "/* language: Alloy */", + "awk" : "// language: AWK", + "thrift" : "/* language: Thrift */", + "sparql" : "# language: SPARQL", + "augeas" : "// language: Augeas", + "cmake" : "# language: CMake", + "f-sharp" : "// language: F#", + "stan" : "// language: Stan", + "isabelle" : "(*language: Isabelle*)", + "dockerfile" : "# language: Dockerfile", + "rmarkdown" : "# language: RMarkdown", + "literate-agda": "-- language: Literate Agda", + "tcl" : "// language: Augeas", + "glsl" : "// language: GLSL", + "antlr" : "// language: ANTLR", + "verilog" : "// language: Verilog", + "racket" : "; language: Racket", + "standard-ml" : "(*language:Standard ML*)", + "elm" : "-- language: Elm", + "yaml" : "# language: YAML", + "smalltalk" : "'' language: Smalltalk", + "ocaml" : "(*language: OCaml*)", + "idris" : "-- language: Idris", + "visual-basic" : "' language: Visual Basic", + "protocol-buffer": "// language: Protocol Buffer", + "bluespec" : "// language: Bluespec", + "applescript" : "-- language: AppleScript", + "makefile" : "# language: Makefile", + "tcsh" : "# language: TCSH", + "maple" : "# language: Maple", + "systemverilog": "// language: SystemVerilog", + "literate-coffeescript": "# language: Literate CoffeeScript", + "vhdl" : "-- language: VHDL", + "restructuredtext": ".. language: reStructuredText", + "sas" : "* language: SAS", + "literate-haskell": "> language: Literate Haskell", + "java-server-pages": "// language: Java Server Pages", + "coffeescript" : "# language: CoffeeScript", + "emacs-lisp" : "; language: Emacs Lisp", + "mathematica" : "// language: Mathematica", + "xslt" : "", + "zig" : "// language: Zig", + "common-lisp" : "; language: Common Lisp", + "stata" : "* language: Stata", + "agda" : "-- language: Agda", + "ada" : "-- language: Ada", +} + + +LANGUAGE_COMMENT_SIGN = {} +for lang in LANGUAGE_TAG: + LANGUAGE_COMMENT_SIGN[lang] = LANGUAGE_TAG[lang].split("language:")[0].strip() + + +IMPORT_HELPER = { + "python": [ + "import math", + "import re", + "import sys", + "import copy", + "import datetime", + "import itertools", + "import collections", + "import heapq", + "import statistics", + "import functools", + "import hashlib", + "import numpy", + "import numpy as np", + "import string", + "from typing import *", + "from collections import *", + ], + "go" : [ + "math", + "strings", + "fmt", + "strconv", + "time", + "bytes", + "regexp", + "sort", + "math/rand", + "crypto/md5", + ], + "cpp" : [ + "#include", + "#include", + "#include", + "#include", + "#include", + "#include", + "#include", + "#include", + "#include", + ], +} + + + +def set_random_seed(seed): + """Set random seed for reproducability.""" + random.seed(seed) + numpy.random.seed(seed) + torch.manual_seed(seed) + + +def stream_jsonl(filename: str) -> Iterable[Dict]: + """ + Parses each jsonl line and yields it as a dictionary + """ + if filename.endswith(".gz"): + with open(filename, "rb") as gzfp: + with gzip.open(gzfp, "rt") as fp: + for line in fp: + if any(not x.isspace() for x in line): + yield json.loads(line) + else: + with open(filename, "r") as fp: + for line in fp: + if any(not x.isspace() for x in line): + yield json.loads(line) + + +def stream_jsonl_all(filename: str) -> Iterable[Dict]: + results = [] + if filename.endswith(".gz"): + fp = gzip.open(open(filename, "rb"), "rt") + else: + fp = open(filename, "r") + for line in fp: + if any(not x.isspace() for x in line): + results.append(json.loads(line)) + fp.close() + + return results + + +def read_dataset( + data_file: str = None, + dataset_type: str = "humanevalx", +) -> Dict: + if "humanevalx" in dataset_type.lower(): + dataset = {task["task_id"]: task for task in stream_jsonl(data_file)} + elif "mbpp" in dataset_type.lower(): + problems = {task["task_id"]: task for task in stream_jsonl(data_file)} + task_ids = sorted(problems.keys())[10:510] + dataset = {} + for task_id in task_ids: + sample = problems[task_id] + description = sample["text"] + test_example = sample["test_list"][0] + prompt = f'"""\n{description}\n{test_example}\n"""\n' + sample["prompt"] = prompt + dataset[task_id] = sample + elif "ds1000" in dataset_type.lower(): + # install ds1000 from https://github.com/HKUNLP/DS-1000 + from ds1000 import DS1000Dataset + ds1000 = DS1000Dataset(source_dir=data_file, libs="all", mode="Completion") + for lib in ds1000.libs: + for problem_id in range(len(ds1000[lib])): + prefix = "" + suffix = "" + insert_flag = False + first_line_flag = True + # extract prefix and suffix of the prompt + for line in ds1000[lib][problem_id]["prompt"].split("\n"): + if "[insert]" in line: + insert_flag = True + continue + if first_line_flag: + first_line_flag = False + else: + line = "\n" + line + if not insert_flag: + prefix += line + else: + suffix += line + + else: + raise f"Dataset: {dataset_type} not supported." + + return dataset + + +def read_translation_dataset( + data_file_src: str = None, + data_file_tgt: str = None, + lang_src: str = None, + lang_tgt: str = None, + dataset_type: str = "humanevalx", +) -> Dict: + if "humanevalx" in dataset_type.lower(): + dataset_src = {task["task_id"]: task for task in stream_jsonl(data_file_src)} + dataset_tgt = {task["task_id"].split("/")[-1]: task for task in stream_jsonl(data_file_tgt)} + for k, sample in dataset_src.items(): + prompt = "code translation\n" + if lang_src == "cpp": + prompt += "C++:\n" + elif lang_src == "js": + prompt += "JavaScript:\n" + else: + prompt += f"{lang_src}:\n".capitalize() + prompt += dataset_src[k]["declaration"] + "\n" + dataset_src[k]["canonical_solution"].rstrip() + "\n" + if lang_tgt == "cpp": + prompt += "C++:\n" + elif lang_tgt == "js": + prompt += "JavaScript:\n" + else: + prompt += f"{lang_tgt}:\n".capitalize() + prompt += dataset_tgt[k.split("/")[-1]]["declaration"] + dataset_src[k]["prompt"] = prompt + else: + raise f"Dataset: {dataset_type} not supported." + + return dataset_src + + +def process_extra_prompt( + prompt: str, + language_type: str = "python", + dataset_type: str = None, + generation_mode: str = "completion", +) -> str: + """ + Processes the extra prompt. + """ + language = language_type.lower() + if dataset_type == "humanevalx": + extra_prompt = "" + # extra_prompt = LANGUAGE_TAG[language] + "\n" + prompt = prompt.strip() + if generation_mode == "instruction": + return "问：" + extra_prompt + prompt + "\n答：" + return extra_prompt + prompt + elif dataset_type == "mbpp": + extra_prompt = "" + prompt = prompt.strip() + return extra_prompt + prompt + else: + return prompt + + +def is_code_generation_finished( + code: str, + dataset_type: str = None, + language_type: str = None, +): + """ + Checks whether the generated code is finished. + """ + if dataset_type == "mbpp": + end_words = ["\ndef", "\nassert"] + for w in end_words: + if w == "\ndef": + if code.count(w) > 1: + return True + else: + if w in code: + return True + else: + if language_type.lower() == "python": + for line in code.split("\n"): + if len(line.strip()) > 0 and line[0] != ' ' and line[0] != '\t': + return True + end_words = ["\ndef", "\nclass", "\nif", "\n#", "\nprint"] + for w in end_words: + if w in code: + return True + elif language_type.lower() == "java": + if code.count("{") + 1 == code.count("}"): + return True + elif language_type.lower() == "go": + if "\nfunc main(" in code: + return True + if code.count("{") + 1 == code.count("}"): + return True + elif language_type.lower() == "js": + if code.count("{") + 1 == code.count("}"): + return True + elif language_type.lower() == "cpp": + if "\nint main()" in code: + return True + if code.count("{") + 1 == code.count("}"): + return True + elif language_type.lower() == "rust": + if "\nfn main()" in code: + return True + if code.count("{") + 1 == code.count("}"): + return True + + return False + + +# Modified from https://github.com/bigcode-project/bigcode-evaluation-harness/blob/main/lm_eval/tasks/mbpp.py +stop_words=["\nclass", "\nassert", '\n"""', "\nprint", "\nif"] +def first_block(string, stop_words): + """Split off first block of code by scanning for class, def etc. on newlines.""" + return re.split("|".join(stop_words), string)[0].rstrip() + + +def cleanup_code( + code: str, + dataset_type: str = None, + language_type: str = None, +): + """ + Cleans up the generated code. + """ + if dataset_type == "mbpp": + end_words = ["\nassert", "\ndef"] + for w in end_words: + if w == "\ndef": + if code.count(w) > 1: + code = code[:code.rfind(w)] + else: + code = code[:code.rfind(w)] + code = first_block(code, stop_words) + elif dataset_type == "humanevalx": + if language_type.lower() == "python": + code_splits = code.split("\n") + is_empty_line = False + ind_empty_line = None + for i, line in enumerate(code_splits): + if len(line.strip()) > 0 and line[0] != ' ' and line[0] != '\t': + is_empty_line = True + ind_empty_line = i + break + if is_empty_line: + code = "\n".join(code_splits[:ind_empty_line]) + else: + end_words = ["\ndef", "\nclass", "\n#", "\nassert", '\n"""', "\nprint", "\nif", "\n\n\n"] + for w in end_words: + if w in code: + code = code[:code.rfind(w)] + elif language_type.lower() == "java": + main_pos = code.find("public static void main") + if main_pos != -1: + code = code[:main_pos] + '}' + if '}' in code: + code = code[:code.rfind('}')] + '}' + if code.count('{') + 1 == code.count('}'): + code += "\n}" + elif language_type.lower() == "go": + if "\nfunc main(" in code: + code = code[:code.rfind("func main(")] + if '}' in code: + code = code[:code.rfind('}')] + '}' + elif language_type.lower() == "cpp": + if "\nint main()" in code: + code = code[:code.rfind("int main()")] + if '}' in code: + code = code[:code.rfind('}')] + '}' + elif language_type.lower() == "js": + if '}' in code: + code = code[:code.rfind('}')] + '}' + elif language_type.lower() == "rust": + if '}' in code: + code = code[:code.rfind('}')] + '}' + + return code + + +def estimate_pass_at_k( + num_samples: Union[int, List[int], np.ndarray], + num_correct: Union[List[int], np.ndarray], + k: int +) -> np.ndarray: + """ + Estimates pass@k of each problem and returns them in an array. + """ + + def estimator(n: int, c: int, k: int) -> float: + """ + Calculates 1 - comb(n - c, k) / comb(n, k). + """ + if n - c < k: + return 1.0 + return 1.0 - np.prod(1.0 - k / np.arange(n - c + 1, n + 1)) + + if isinstance(num_samples, int): + num_samples_it = itertools.repeat(num_samples, len(num_correct)) + else: + assert len(num_samples) == len(num_correct) + num_samples_it = iter(num_samples) + + return np.array([estimator(int(n), int(c), k) for n, c in zip(num_samples_it, num_correct)]) + + +class Logger: + def __init__(self, name, log_level=logging.INFO, log_file=None, log_mode="both", disable_formatter=False): + self.logger = logging.getLogger(name) + self.logger.setLevel(log_level) + + self.formatter = logging.Formatter('%(asctime)s [%(levelname)s] %(message)s', datefmt='%Y-%m-%d %H:%M:%S') + + # Log to console + if log_mode == "both" or log_mode == "terminal": + console_handler = logging.StreamHandler() + if not disable_formatter: + console_handler.setFormatter(self.formatter) + self.logger.addHandler(console_handler) + + # Log to file + if log_file is not None: + if log_mode == "both" or log_mode == "file": + file_handler = logging.FileHandler(log_file, mode='w') + if not disable_formatter: + file_handler.setFormatter(self.formatter) + self.logger.addHandler(file_handler) + + def add_file_handler(self, file_name): + file_handler = logging.FileHandler(file_name, mode='w') + file_handler.setFormatter(self.formatter) + self.logger.addHandler(file_handler) + + def debug(self, message): + self.logger.debug(message) + + def info(self, message): + self.logger.info(message) + + def warning(self, message): + self.logger.warning(message) + + def error(self, message): + self.logger.error(message) + + def critical(self, message): + self.logger.critical(message) diff --git a/PyTorch/built-in/foundation/CodeGeeX2/fix/modeling_utils.py b/PyTorch/built-in/foundation/CodeGeeX2/fix/modeling_utils.py new file mode 100644 index 0000000000000000000000000000000000000000..a8f80d6f235e4d15ec5aa908f364ae28c2a47d4b --- /dev/null +++ b/PyTorch/built-in/foundation/CodeGeeX2/fix/modeling_utils.py @@ -0,0 +1,3943 @@ +# coding=utf-8 +# Copyright 2023 Huawei Technologies Co., Ltd +# +# 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. + +# coding=utf-8 +# Copyright 2018 The Google AI Language Team Authors, Facebook AI Research authors and The HuggingFace Inc. team. +# Copyright (c) 2018, NVIDIA CORPORATION. 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 collections +import gc +import inspect +import json +import os +import re +import shutil +import tempfile +import warnings +from contextlib import contextmanager +from dataclasses import dataclass +from functools import partial +from typing import Any, Callable, Dict, List, Optional, Tuple, Union + +import torch +from packaging import version +from torch import Tensor, nn +from torch.nn import CrossEntropyLoss + +from .activations import get_activation +from .configuration_utils import PretrainedConfig +from .deepspeed import deepspeed_config, is_deepspeed_zero3_enabled +from .dynamic_module_utils import custom_object_save +from .generation import GenerationConfig, GenerationMixin +from .pytorch_utils import ( # noqa: F401 + Conv1D, + apply_chunking_to_forward, + find_pruneable_heads_and_indices, + id_tensor_storage, + prune_conv1d_layer, + prune_layer, + prune_linear_layer, +) +from .utils import ( + DUMMY_INPUTS, + FLAX_WEIGHTS_NAME, + SAFE_WEIGHTS_INDEX_NAME, + SAFE_WEIGHTS_NAME, + TF2_WEIGHTS_NAME, + TF_WEIGHTS_NAME, + WEIGHTS_INDEX_NAME, + WEIGHTS_NAME, + ContextManagers, + ModelOutput, + PushToHubMixin, + cached_file, + copy_func, + download_url, + has_file, + is_accelerate_available, + is_bitsandbytes_available, + is_offline_mode, + is_optimum_available, + is_remote_url, + is_safetensors_available, + is_torch_tpu_available, + logging, + replace_return_docstrings, +) +from .utils.hub import convert_file_size_to_int, get_checkpoint_shard_files +from .utils.import_utils import ENV_VARS_TRUE_VALUES, importlib_metadata, is_sagemaker_mp_enabled +from .utils.quantization_config import BitsAndBytesConfig +from .utils.versions import require_version_core + +XLA_USE_BF16 = os.environ.get("XLA_USE_BF16", "0").upper() +XLA_DOWNCAST_BF16 = os.environ.get("XLA_DOWNCAST_BF16", "0").upper() + +if is_accelerate_available(): + from accelerate import __version__ as accelerate_version + from accelerate import dispatch_model, infer_auto_device_map, init_empty_weights + from accelerate.utils import ( + find_tied_parameters, + load_offloaded_weights, + offload_weight, + save_offload_index, + set_module_tensor_to_device, + ) + + if version.parse(accelerate_version) > version.parse("0.11.0"): + from accelerate.utils import get_balanced_memory + else: + get_balanced_memory = None + if version.parse(accelerate_version) > version.parse("0.19.0"): + from accelerate.utils import check_tied_parameters_on_same_device + else: + check_tied_parameters_on_same_device = None +else: + find_tied_parameters = None + +if is_safetensors_available(): + from safetensors import safe_open + from safetensors.torch import load_file as safe_load_file + from safetensors.torch import save_file as safe_save_file + +logger = logging.get_logger(__name__) + +_init_weights = True + +if is_sagemaker_mp_enabled(): + import smdistributed.modelparallel.torch as smp + from smdistributed.modelparallel import __version__ as SMP_VERSION + + IS_SAGEMAKER_MP_POST_1_10 = version.parse(SMP_VERSION) >= version.parse("1.10") +else: + IS_SAGEMAKER_MP_POST_1_10 = False + + +@contextmanager +def no_init_weights(_enable=True): + """ + Context manager to globally disable weight initialization to speed up loading large models. + + TODO(Patrick): Delete safety argument `_enable=True` at next major version. . + """ + global _init_weights + old_init_weights = _init_weights + if _enable: + _init_weights = False + try: + yield + finally: + _init_weights = old_init_weights + + +try: + from torch.nn import Identity +except ImportError: + # Older PyTorch compatibility + class Identity(nn.Module): + r"""A placeholder identity operator that is argument-insensitive.""" + + def __init__(self, *args, **kwargs): + super().__init__() + + def forward(self, input): + return input + + +def get_parameter_device(parameter: Union[nn.Module, GenerationMixin, "ModuleUtilsMixin"]): + try: + return next(parameter.parameters()).device + except StopIteration: + # For nn.DataParallel compatibility in PyTorch 1.5 + + def find_tensor_attributes(module: nn.Module) -> List[Tuple[str, Tensor]]: + tuples = [(k, v) for k, v in module.__dict__.items() if torch.is_tensor(v)] + return tuples + + gen = parameter._named_members(get_members_fn=find_tensor_attributes) + first_tuple = next(gen) + return first_tuple[1].device + + +def get_first_parameter_dtype(parameter: Union[nn.Module, GenerationMixin, "ModuleUtilsMixin"]): + """ + Returns the first parameter dtype (can be non-floating) or asserts if none were found. + """ + try: + return next(parameter.parameters()).dtype + except StopIteration: + # For nn.DataParallel compatibility in PyTorch > 1.5 + + def find_tensor_attributes(module: nn.Module) -> List[Tuple[str, Tensor]]: + tuples = [(k, v) for k, v in module.__dict__.items() if torch.is_tensor(v)] + return tuples + + gen = parameter._named_members(get_members_fn=find_tensor_attributes) + first_tuple = next(gen) + return first_tuple[1].dtype + + +def get_parameter_dtype(parameter: Union[nn.Module, GenerationMixin, "ModuleUtilsMixin"]): + """ + Returns the first found floating dtype in parameters if there is one, otherwise returns the last dtype it found. + """ + last_dtype = None + for t in parameter.parameters(): + last_dtype = t.dtype + if t.is_floating_point(): + # Adding fix for https://github.com/pytorch/xla/issues/4152 + # Fixes issue where the model code passes a value that is out of range for XLA_USE_BF16=1 + # and XLA_DOWNCAST_BF16=1 so the conversion would cast it to -inf + # NOTE: `is_torch_tpu_available()` is checked last as it induces a graph break in torch dynamo + if XLA_USE_BF16 in ENV_VARS_TRUE_VALUES and is_torch_tpu_available(): + return torch.bfloat16 + if XLA_DOWNCAST_BF16 in ENV_VARS_TRUE_VALUES and is_torch_tpu_available(): + if t.dtype == torch.float: + return torch.bfloat16 + if t.dtype == torch.double: + return torch.float32 + return t.dtype + + if last_dtype is not None: + # if no floating dtype was found return whatever the first dtype is + return last_dtype + + # For nn.DataParallel compatibility in PyTorch > 1.5 + def find_tensor_attributes(module: nn.Module) -> List[Tuple[str, Tensor]]: + tuples = [(k, v) for k, v in module.__dict__.items() if torch.is_tensor(v)] + return tuples + + gen = parameter._named_members(get_members_fn=find_tensor_attributes) + last_tuple = None + for tuple in gen: + last_tuple = tuple + if tuple[1].is_floating_point(): + return tuple[1].dtype + + if last_tuple is not None: + # fallback to the last dtype + return last_tuple[1].dtype + + # fallback to buffer dtype + for t in parameter.buffers(): + last_dtype = t.dtype + if t.is_floating_point(): + return t.dtype + return last_dtype + + +def get_state_dict_float_dtype(state_dict): + """ + Returns the first found floating dtype in `state_dict` or asserts if none were found. + """ + for t in state_dict.values(): + if t.is_floating_point(): + return t.dtype + + raise ValueError("couldn't find any floating point dtypes in state_dict") + + +def get_state_dict_dtype(state_dict): + """ + Returns the first found floating dtype in `state_dict` if there is one, otherwise returns the first dtype. + """ + for t in state_dict.values(): + if t.is_floating_point(): + return t.dtype + + # if no floating dtype was found return whatever the first dtype is + else: + return next(state_dict.values()).dtype + + +def dtype_byte_size(dtype): + """ + Returns the size (in bytes) occupied by one parameter of type `dtype`. + + Example: + + ```py + >>> dtype_byte_size(torch.float32) + 4 + ``` + """ + if dtype == torch.bool: + return 1 / 8 + bit_search = re.search(r"[^\d](\d+)$", str(dtype)) + if bit_search is None: + raise ValueError(f"`dtype` is not a valid dtype: {dtype}.") + bit_size = int(bit_search.groups()[0]) + return bit_size // 8 + + +def shard_checkpoint( + state_dict: Dict[str, torch.Tensor], max_shard_size: Union[int, str] = "10GB", weights_name: str = WEIGHTS_NAME +): + """ + Splits a model state dictionary in sub-checkpoints so that the final size of each sub-checkpoint does not exceed a + given size. + + The sub-checkpoints are determined by iterating through the `state_dict` in the order of its keys, so there is no + optimization made to make each sub-checkpoint as close as possible to the maximum size passed. For example, if the + limit is 10GB and we have weights of sizes [6GB, 6GB, 2GB, 6GB, 2GB, 2GB] they will get sharded as [6GB], [6+2GB], + [6+2+2GB] and not [6+2+2GB], [6+2GB], [6GB]. + + + + If one of the model's weight is bigger that `max_sahrd_size`, it will end up in its own sub-checkpoint which will + have a size greater than `max_shard_size`. + + + + Args: + state_dict (`Dict[str, torch.Tensor]`): The state dictionary of a model to save. + max_shard_size (`int` or `str`, *optional*, defaults to `"10GB"`): + The maximum size of each sub-checkpoint. If expressed as a string, needs to be digits followed by a unit + (like `"5MB"`). + weights_name (`str`, *optional*, defaults to `"pytorch_model.bin"`): + The name of the model save file. + """ + max_shard_size = convert_file_size_to_int(max_shard_size) + ''' + sharded_state_dicts = [{}] + last_block_size = 0 + total_size = 0 + storage_id_to_block = {} + + for key, weight in state_dict.items(): + storage_id = id_tensor_storage(weight) + + # If a `weight` shares the same underlying storage as another tensor, we put `weight` in the same `block` + if storage_id in storage_id_to_block: + block_id = storage_id_to_block[storage_id] + sharded_state_dicts[block_id][key] = weight + continue + + weight_size = weight.numel() * dtype_byte_size(weight.dtype) + + # If this weight is going to tip up over the maximal size, we split. + if last_block_size + weight_size > max_shard_size: + sharded_state_dicts.append({}) + last_block_size = 0 + + sharded_state_dicts[-1][key] = weight + last_block_size += weight_size + total_size += weight_size + storage_id_to_block[storage_id] = len(sharded_state_dicts) - 1 + ''' + sharded_state_dicts = [] + current_block = {} + current_block_size = 0 + total_size = 0 + + for key, weight in state_dict.items(): + weight_size = weight.numel() * dtype_byte_size(weight.dtype) + + # If this weight is going to tip up over the maximal size, we split. + if current_block_size + weight_size > max_shard_size: + sharded_state_dicts.append(current_block) + current_block = {} + current_block_size = 0 + + current_block[key] = weight + current_block_size += weight_size + total_size += weight_size + + # Add the last block + sharded_state_dicts.append(current_block) + # If we only have one shard, we return it + if len(sharded_state_dicts) == 1: + return {weights_name: sharded_state_dicts[0]}, None + + # Otherwise, let's build the index + weight_map = {} + shards = {} + for idx, shard in enumerate(sharded_state_dicts): + shard_file = weights_name.replace(".bin", f"-{idx + 1:05d}-of-{len(sharded_state_dicts):05d}.bin") + shard_file = shard_file.replace( + ".safetensors", f"-{idx + 1:05d}-of-{len(sharded_state_dicts):05d}.safetensors" + ) + shards[shard_file] = shard + for key in shard.keys(): + weight_map[key] = shard_file + + # Add the metadata + metadata = {"total_size": total_size} + index = {"metadata": metadata, "weight_map": weight_map} + return shards, index + + +def load_sharded_checkpoint(model, folder, strict=True, prefer_safe=True): + """ + This is the same as + [`torch.nn.Module.load_state_dict`](https://pytorch.org/docs/stable/generated/torch.nn.Module.html?highlight=load_state_dict#torch.nn.Module.load_state_dict) + but for a sharded checkpoint. + + This load is performed efficiently: each checkpoint shard is loaded one by one in RAM and deleted after being + loaded in the model. + + Args: + model (`torch.nn.Module`): The model in which to load the checkpoint. + folder (`str` or `os.PathLike`): A path to a folder containing the sharded checkpoint. + strict (`bool`, *optional`, defaults to `True`): + Whether to strictly enforce that the keys in the model state dict match the keys in the sharded checkpoint. + prefer_safe (`bool`, *optional*, defaults to `False`) + If both safetensors and PyTorch save files are present in checkpoint and `prefer_safe` is True, the + safetensors files will be loaded. Otherwise, PyTorch files are always loaded when possible. + + Returns: + `NamedTuple`: A named tuple with `missing_keys` and `unexpected_keys` fields + - `missing_keys` is a list of str containing the missing keys + - `unexpected_keys` is a list of str containing the unexpected keys + """ + # Load the index + index_file = os.path.join(folder, WEIGHTS_INDEX_NAME) + safe_index_file = os.path.join(folder, SAFE_WEIGHTS_INDEX_NAME) + + index_present = os.path.isfile(index_file) + safe_index_present = os.path.isfile(safe_index_file) + + if not index_present and not (safe_index_present and is_safetensors_available()): + filenames = ( + (WEIGHTS_INDEX_NAME, SAFE_WEIGHTS_INDEX_NAME) if is_safetensors_available() else (WEIGHTS_INDEX_NAME,) + ) + raise ValueError(f"Can't find a checkpoint index ({' or '.join(filenames)}) in {folder}.") + + load_safe = False + if safe_index_present: + if prefer_safe: + if is_safetensors_available(): + load_safe = True # load safe due to preference + else: + logger.warning( + f"Cannot load sharded checkpoint at {folder} safely since safetensors is not installed!" + ) + elif not index_present: + load_safe = True # load safe since we have no other choice + + load_index = safe_index_file if load_safe else index_file + + with open(load_index, "r", encoding="utf-8") as f: + index = json.load(f) + + shard_files = list(set(index["weight_map"].values())) + + # If strict=True, error before loading any of the state dicts. + loaded_keys = index["weight_map"].keys() + model_keys = model.state_dict().keys() + missing_keys = [key for key in model_keys if key not in loaded_keys] + unexpected_keys = [key for key in loaded_keys if key not in model_keys] + if strict and (len(missing_keys) > 0 or len(unexpected_keys) > 0): + error_message = f"Error(s) in loading state_dict for {model.__class__.__name__}" + if len(missing_keys) > 0: + str_missing_keys = ",".join([f'"{k}"' for k in missing_keys]) + error_message += f"\nMissing key(s): {str_missing_keys}." + if len(unexpected_keys) > 0: + str_unexpected_keys = ",".join([f'"{k}"' for k in unexpected_keys]) + error_message += f"\nMissing key(s): {str_unexpected_keys}." + raise RuntimeError(error_message) + + loader = safe_load_file if load_safe else partial(torch.load, map_location="cpu") + + for shard_file in shard_files: + state_dict = loader(os.path.join(folder, shard_file)) + model.load_state_dict(state_dict, strict=False) + + # Make sure memory is freed before we load the next state dict. + del state_dict + gc.collect() + + # Return the same thing as PyTorch load_state_dict function. + return torch.nn.modules.module._IncompatibleKeys(missing_keys, unexpected_keys) + + +def load_state_dict(checkpoint_file: Union[str, os.PathLike]): + """ + Reads a PyTorch checkpoint file, returning properly formatted errors if they arise. + """ + if checkpoint_file.endswith(".safetensors") and is_safetensors_available(): + # Check format of the archive + with safe_open(checkpoint_file, framework="pt") as f: + metadata = f.metadata() + if metadata.get("format") not in ["pt", "tf", "flax"]: + raise OSError( + f"The safetensors archive passed at {checkpoint_file} does not contain the valid metadata. Make sure " + "you save your model with the `save_pretrained` method." + ) + elif metadata["format"] != "pt": + raise NotImplementedError( + f"Conversion from a {metadata['format']} safetensors archive to PyTorch is not implemented yet." + ) + return safe_load_file(checkpoint_file) + try: + return torch.load(checkpoint_file, map_location="cpu") + except Exception as e: + try: + with open(checkpoint_file) as f: + if f.read(7) == "version": + raise OSError( + "You seem to have cloned a repository without having git-lfs installed. Please install " + "git-lfs and run `git lfs install` followed by `git lfs pull` in the folder " + "you cloned." + ) + else: + raise ValueError( + f"Unable to locate the file {checkpoint_file} which is necessary to load this pretrained " + "model. Make sure you have saved the model properly." + ) from e + except (UnicodeDecodeError, ValueError): + raise OSError( + f"Unable to load weights from pytorch checkpoint file for '{checkpoint_file}' " + f"at '{checkpoint_file}'. " + "If you tried to load a PyTorch model from a TF 2.0 checkpoint, please set from_tf=True." + ) + + +def set_initialized_submodules(model, state_dict_keys): + """ + Sets the `_is_hf_initialized` flag in all submodules of a given model when all its weights are in the loaded state + dict. + """ + for module_name, module in model.named_modules(): + loaded_keys = [k.replace(f"{module_name}.", "") for k in state_dict_keys if k.startswith(f"{module_name}.")] + if len(set(module.state_dict().keys()) - set(loaded_keys)) == 0: + module._is_hf_initialized = True + + +def _load_state_dict_into_model(model_to_load, state_dict, start_prefix): + # Convert old format to new format if needed from a PyTorch state_dict + old_keys = [] + new_keys = [] + for key in state_dict.keys(): + new_key = None + if "gamma" in key: + new_key = key.replace("gamma", "weight") + if "beta" in key: + new_key = key.replace("beta", "bias") + if new_key: + old_keys.append(key) + new_keys.append(new_key) + for old_key, new_key in zip(old_keys, new_keys): + state_dict[new_key] = state_dict.pop(old_key) + + # copy state_dict so _load_from_state_dict can modify it + metadata = getattr(state_dict, "_metadata", None) + state_dict = state_dict.copy() + if metadata is not None: + state_dict._metadata = metadata + + error_msgs = [] + + # PyTorch's `_load_from_state_dict` does not copy parameters in a module's descendants + # so we need to apply the function recursively. + def load(module: nn.Module, state_dict, prefix=""): + local_metadata = {} if metadata is None else metadata.get(prefix[:-1], {}) + args = (state_dict, prefix, local_metadata, True, [], [], error_msgs) + # Parameters of module and children will start with prefix. We can exit early if there are none in this + # state_dict + if len([key for key in state_dict if key.startswith(prefix)]) > 0: + if is_deepspeed_zero3_enabled(): + import deepspeed + + # In sharded models, each shard has only part of the full state_dict, so only gather + # parameters that are in the current state_dict. + named_parameters = dict(module.named_parameters(prefix=prefix[:-1], recurse=False)) + params_to_gather = [named_parameters[k] for k in state_dict.keys() if k in named_parameters] + if len(params_to_gather) > 0: + # because zero3 puts placeholders in model params, this context + # manager gathers (unpartitions) the params of the current layer, then loads from + # the state dict and then re-partitions them again + with deepspeed.zero.GatheredParameters(params_to_gather, modifier_rank=0): + if torch.distributed.get_rank() == 0: + module._load_from_state_dict(*args) + else: + module._load_from_state_dict(*args) + + for name, child in module._modules.items(): + if child is not None: + load(child, state_dict, prefix + name + ".") + + load(model_to_load, state_dict, prefix=start_prefix) + # Delete `state_dict` so it could be collected by GC earlier. Note that `state_dict` is a copy of the argument, so + # it's safe to delete it. + del state_dict + + return error_msgs + + +def find_submodule_and_param_name(model, long_key, start_prefix): + """ + A helper util to find the last sub-module and the param/buffer name. If `start_prefix` is supplied it'll be removed + from the start of the key + """ + + if len(start_prefix) > 0 and long_key.startswith(start_prefix): + long_key = ".".join(long_key.split(".")[1:]) + + split_key = long_key.split(".") + submodule = model + while len(split_key) > 1: + if hasattr(submodule, split_key[0]): + submodule = getattr(submodule, split_key[0]) + del split_key[0] + else: + submodule = None + break + if submodule == model: + submodule = None + return submodule, split_key[0] + + +def _move_model_to_meta(model, loaded_state_dict_keys, start_prefix): + """ + Moves `loaded_state_dict_keys` in model to meta device which frees up the memory taken by those params. + + `start_prefix` is used for models which insert their name into model keys, e.g. `bert` in + `bert.pooler.dense.weight` + + """ + + # dematerialize param storage for keys that are going to be replaced by state_dict, by + # putting those on the meta device + for k in loaded_state_dict_keys: + submodule, param_name = find_submodule_and_param_name(model, k, start_prefix) + if submodule is not None: + # selectively switch to the meta device only those params/buffers that will + # be next replaced from state_dict. This a complex way to do p.to_("meta") + # since we have no in-place to_ for tensors. + new_val = getattr(submodule, param_name) + if isinstance(new_val, torch.nn.Parameter): + # isinstance returns False for Params on meta device, so switch after the check + new_val = torch.nn.Parameter(new_val.to("meta")) + else: + new_val = new_val.to("meta") + setattr(submodule, param_name, new_val) + + +def _load_state_dict_into_meta_model( + model, + state_dict, + loaded_state_dict_keys, # left for now but could be removed, see below + start_prefix, + expected_keys, + device_map=None, + offload_folder=None, + offload_index=None, + state_dict_folder=None, + state_dict_index=None, + dtype=None, + is_quantized=False, + is_safetensors=False, + keep_in_fp32_modules=None, +): + """ + This is somewhat similar to `_load_state_dict_into_model`, but deals with a model that has some or all of its + params on a `meta` device. It replaces the model params with the data from the `state_dict`, while moving the + params back to the normal device, but only for `loaded_state_dict_keys`. + + `start_prefix` is used for models which insert their name into model keys, e.g. `bert` in + `bert.pooler.dense.weight` + + """ + + # XXX: remaining features to implement to be fully compatible with _load_state_dict_into_model + # - deepspeed zero 3 support + # - need to copy metadata if any - see _load_state_dict_into_model + # - handling error_msgs - mimicking the error handling in module._load_from_state_dict() + # - Is there a situation where some keys aren't in `loaded_state_dict_keys` and in which case + # they won't get loaded. + + if is_quantized: + from .utils.bitsandbytes import set_module_quantized_tensor_to_device + + error_msgs = [] + + old_keys = [] + new_keys = [] + for key in state_dict.keys(): + new_key = None + if "gamma" in key: + new_key = key.replace("gamma", "weight") + if "beta" in key: + new_key = key.replace("beta", "bias") + if new_key: + old_keys.append(key) + new_keys.append(new_key) + for old_key, new_key in zip(old_keys, new_keys): + state_dict[new_key] = state_dict.pop(old_key) + + for param_name, param in state_dict.items(): + # First part of the test is always true as load_state_dict_keys always contains state_dict keys. + if param_name not in loaded_state_dict_keys or param_name not in expected_keys: + continue + + if param_name.startswith(start_prefix): + param_name = param_name[len(start_prefix):] + + module_name = param_name + set_module_kwargs = {} + + # We convert floating dtypes to the `dtype` passed. We want to keep the buffers/params + # in int/uint/bool and not cast them. + if dtype is not None and torch.is_floating_point(param): + if ( + keep_in_fp32_modules is not None + and any(module_to_keep_in_fp32 in param_name for module_to_keep_in_fp32 in keep_in_fp32_modules) + and dtype == torch.float16 + ): + param = param.to(torch.float32) + + # For backward compatibility with older versions of `accelerate` + # TODO: @sgugger replace this check with version check at the next `accelerate` release + if "dtype" in list(inspect.signature(set_module_tensor_to_device).parameters): + set_module_kwargs["dtype"] = torch.float32 + else: + param = param.to(dtype) + + # For compatibility with PyTorch load_state_dict which converts state dict dtype to existing dtype in model + if dtype is None: + old_param = model + splits = param_name.split(".") + for split in splits: + old_param = getattr(old_param, split) + if old_param is None: + break + + if old_param is not None: + param = param.to(old_param.dtype) + + set_module_kwargs["value"] = param + + if device_map is None: + param_device = "cpu" + else: + # find next higher level module that is defined in device_map: + # bert.lm_head.weight -> bert.lm_head -> bert -> '' + while len(module_name) > 0 and module_name not in device_map: + module_name = ".".join(module_name.split(".")[:-1]) + if module_name == "" and "" not in device_map: + # TODO: group all errors and raise at the end. + raise ValueError(f"{param_name} doesn't have any device set.") + param_device = device_map[module_name] + + if param_device == "disk": + if not is_safetensors: + offload_index = offload_weight(param, param_name, offload_folder, offload_index) + elif param_device == "cpu" and state_dict_index is not None: + state_dict_index = offload_weight(param, param_name, state_dict_folder, state_dict_index) + elif not is_quantized: + # For backward compatibility with older versions of `accelerate` + set_module_tensor_to_device(model, param_name, param_device, **set_module_kwargs) + else: + if param.dtype == torch.int8 and param_name.replace("weight", "SCB") in state_dict.keys(): + fp16_statistics = state_dict[param_name.replace("weight", "SCB")] + else: + fp16_statistics = None + + if "SCB" not in param_name: + set_module_quantized_tensor_to_device( + model, param_name, param_device, value=param, fp16_statistics=fp16_statistics + ) + + return error_msgs, offload_index, state_dict_index + + +def _add_variant(weights_name: str, variant: Optional[str] = None) -> str: + if variant is not None: + splits = weights_name.split(".") + splits = splits[:-1] + [variant] + splits[-1:] + weights_name = ".".join(splits) + + return weights_name + + +class ModuleUtilsMixin: + """ + A few utilities for `torch.nn.Modules`, to be used as a mixin. + """ + + @staticmethod + def _hook_rss_memory_pre_forward(module, *args, **kwargs): + try: + import psutil + except ImportError: + raise ImportError("You need to install psutil (pip install psutil) to use memory tracing.") + + process = psutil.Process(os.getpid()) + mem = process.memory_info() + module.mem_rss_pre_forward = mem.rss + return None + + @staticmethod + def _hook_rss_memory_post_forward(module, *args, **kwargs): + try: + import psutil + except ImportError: + raise ImportError("You need to install psutil (pip install psutil) to use memory tracing.") + + process = psutil.Process(os.getpid()) + mem = process.memory_info() + module.mem_rss_post_forward = mem.rss + mem_rss_diff = module.mem_rss_post_forward - module.mem_rss_pre_forward + module.mem_rss_diff = mem_rss_diff + (module.mem_rss_diff if hasattr(module, "mem_rss_diff") else 0) + return None + + def add_memory_hooks(self): + """ + Add a memory hook before and after each sub-module forward pass to record increase in memory consumption. + + Increase in memory consumption is stored in a `mem_rss_diff` attribute for each module and can be reset to zero + with `model.reset_memory_hooks_state()`. + """ + for module in self.modules(): + module.register_forward_pre_hook(self._hook_rss_memory_pre_forward) + module.register_forward_hook(self._hook_rss_memory_post_forward) + self.reset_memory_hooks_state() + + def reset_memory_hooks_state(self): + """ + Reset the `mem_rss_diff` attribute of each module (see [`~modeling_utils.ModuleUtilsMixin.add_memory_hooks`]). + """ + for module in self.modules(): + module.mem_rss_diff = 0 + module.mem_rss_post_forward = 0 + module.mem_rss_pre_forward = 0 + + @property + def device(self) -> torch.device: + """ + `torch.device`: The device on which the module is (assuming that all the module parameters are on the same + device). + """ + return get_parameter_device(self) + + @property + def dtype(self) -> torch.dtype: + """ + `torch.dtype`: The dtype of the module (assuming that all the module parameters have the same dtype). + """ + return get_parameter_dtype(self) + + def invert_attention_mask(self, encoder_attention_mask: Tensor) -> Tensor: + """ + Invert an attention mask (e.g., switches 0. and 1.). + + Args: + encoder_attention_mask (`torch.Tensor`): An attention mask. + + Returns: + `torch.Tensor`: The inverted attention mask. + """ + if encoder_attention_mask.dim() == 3: + encoder_extended_attention_mask = encoder_attention_mask[:, None, :, :] + if encoder_attention_mask.dim() == 2: + encoder_extended_attention_mask = encoder_attention_mask[:, None, None, :] + # T5 has a mask that can compare sequence ids, we can simulate this here with this transposition + # Cf. https://github.com/tensorflow/mesh/blob/8d2465e9bc93129b913b5ccc6a59aa97abd96ec6/mesh_tensorflow + # /transformer/transformer_layers.py#L270 + # encoder_extended_attention_mask = (encoder_extended_attention_mask == + # encoder_extended_attention_mask.transpose(-1, -2)) + encoder_extended_attention_mask = encoder_extended_attention_mask.to(dtype=self.dtype) # fp16 compatibility + encoder_extended_attention_mask = (1.0 - encoder_extended_attention_mask) * torch.finfo(self.dtype).min + + return encoder_extended_attention_mask + + @staticmethod + def create_extended_attention_mask_for_decoder(input_shape, attention_mask, device=None): + if device is not None: + warnings.warn( + "The `device` argument is deprecated and will be removed in v5 of Transformers.", FutureWarning + ) + else: + device = attention_mask.device + batch_size, seq_length = input_shape + seq_ids = torch.arange(seq_length, device=device) + causal_mask = seq_ids[None, None, :].repeat(batch_size, seq_length, 1) <= seq_ids[None, :, None] + # in case past_key_values are used we need to add a prefix ones mask to the causal mask + # causal and attention masks must have same type with pytorch version < 1.3 + causal_mask = causal_mask.to(attention_mask.dtype) + + if causal_mask.shape[1] < attention_mask.shape[1]: + prefix_seq_len = attention_mask.shape[1] - causal_mask.shape[1] + causal_mask = torch.cat( + [ + torch.ones((batch_size, seq_length, prefix_seq_len), device=device, dtype=causal_mask.dtype), + causal_mask, + ], + axis=-1, + ) + + extended_attention_mask = causal_mask[:, None, :, :] * attention_mask[:, None, None, :] + return extended_attention_mask + + def get_extended_attention_mask( + self, attention_mask: Tensor, input_shape: Tuple[int], device: torch.device = None, + dtype: torch.float = None + ) -> Tensor: + """ + Makes broadcastable attention and causal masks so that future and masked tokens are ignored. + + Arguments: + attention_mask (`torch.Tensor`): + Mask with ones indicating tokens to attend to, zeros for tokens to ignore. + input_shape (`Tuple[int]`): + The shape of the input to the model. + + Returns: + `torch.Tensor` The extended attention mask, with a the same dtype as `attention_mask.dtype`. + """ + if dtype is None: + dtype = self.dtype + + if not (attention_mask.dim() == 2 and self.config.is_decoder): + # show warning only if it won't be shown in `create_extended_attention_mask_for_decoder` + if device is not None: + warnings.warn( + "The `device` argument is deprecated and will be removed in v5 of Transformers.", FutureWarning + ) + # We can provide a self-attention mask of dimensions [batch_size, from_seq_length, to_seq_length] + # ourselves in which case we just need to make it broadcastable to all heads. + if attention_mask.dim() == 3: + extended_attention_mask = attention_mask[:, None, :, :] + elif attention_mask.dim() == 2: + # Provided a padding mask of dimensions [batch_size, seq_length] + # - if the model is a decoder, apply a causal mask in addition to the padding mask + # - if the model is an encoder, make the mask broadcastable to [batch_size, num_heads, seq_length, seq_length] + if self.config.is_decoder: + extended_attention_mask = ModuleUtilsMixin.create_extended_attention_mask_for_decoder( + input_shape, attention_mask, device + ) + else: + extended_attention_mask = attention_mask[:, None, None, :] + else: + raise ValueError( + f"Wrong shape for input_ids (shape {input_shape}) or attention_mask (shape {attention_mask.shape})" + ) + + # Since attention_mask is 1.0 for positions we want to attend and 0.0 for + # masked positions, this operation will create a tensor which is 0.0 for + # positions we want to attend and the dtype's smallest value for masked positions. + # Since we are adding it to the raw scores before the softmax, this is + # effectively the same as removing these entirely. + extended_attention_mask = extended_attention_mask.to(dtype=dtype) # fp16 compatibility + extended_attention_mask = (1.0 - extended_attention_mask) * torch.finfo(dtype).min + return extended_attention_mask + + def get_head_mask( + self, head_mask: Optional[Tensor], num_hidden_layers: int, is_attention_chunked: bool = False + ) -> Tensor: + """ + Prepare the head mask if needed. + + Args: + head_mask (`torch.Tensor` with shape `[num_heads]` or `[num_hidden_layers x num_heads]`, *optional*): + The mask indicating if we should keep the heads or not (1.0 for keep, 0.0 for discard). + num_hidden_layers (`int`): + The number of hidden layers in the model. + is_attention_chunked (`bool`, *optional*, defaults to `False`): + Whether or not the attentions scores are computed by chunks or not. + + Returns: + `torch.Tensor` with shape `[num_hidden_layers x batch x num_heads x seq_length x seq_length]` or list with + `[None]` for each layer. + """ + if head_mask is not None: + head_mask = self._convert_head_mask_to_5d(head_mask, num_hidden_layers) + if is_attention_chunked is True: + head_mask = head_mask.unsqueeze(-1) + else: + head_mask = [None] * num_hidden_layers + + return head_mask + + def _convert_head_mask_to_5d(self, head_mask, num_hidden_layers): + """-> [num_hidden_layers x batch x num_heads x seq_length x seq_length]""" + if head_mask.dim() == 1: + head_mask = head_mask.unsqueeze(0).unsqueeze(0).unsqueeze(-1).unsqueeze(-1) + head_mask = head_mask.expand(num_hidden_layers, -1, -1, -1, -1) + elif head_mask.dim() == 2: + head_mask = head_mask.unsqueeze(1).unsqueeze(-1).unsqueeze(-1) # We can specify head_mask for each layer + assert head_mask.dim() == 5, f"head_mask.dim != 5, instead {head_mask.dim()}" + head_mask = head_mask.to(dtype=self.dtype) # switch to float if need + fp16 compatibility + return head_mask + + def num_parameters(self, only_trainable: bool = False, exclude_embeddings: bool = False) -> int: + """ + Get number of (optionally, trainable or non-embeddings) parameters in the module. + + Args: + only_trainable (`bool`, *optional*, defaults to `False`): + Whether or not to return only the number of trainable parameters + + exclude_embeddings (`bool`, *optional*, defaults to `False`): + Whether or not to return only the number of non-embeddings parameters + + Returns: + `int`: The number of parameters. + """ + + if exclude_embeddings: + embedding_param_names = [ + f"{name}.weight" for name, module_type in self.named_modules() if isinstance(module_type, nn.Embedding) + ] + non_embedding_parameters = [ + parameter for name, parameter in self.named_parameters() if name not in embedding_param_names + ] + return sum(p.numel() for p in non_embedding_parameters if p.requires_grad or not only_trainable) + else: + return sum(p.numel() for p in self.parameters() if p.requires_grad or not only_trainable) + + def estimate_tokens(self, input_dict: Dict[str, Union[torch.Tensor, Any]]) -> int: + """ + Helper function to estimate the total number of tokens from the model inputs. + + Args: + inputs (`dict`): The model inputs. + + Returns: + `int`: The total number of tokens. + """ + if not hasattr(self, "warnings_issued"): + self.warnings_issued = {} + if self.main_input_name in input_dict: + return input_dict[self.main_input_name].numel() + elif "estimate_tokens" not in self.warnings_issued: + logger.warning( + "Could not estimate the number of tokens of the input, floating-point operations will not be computed" + ) + self.warnings_issued["estimate_tokens"] = True + return 0 + + def floating_point_ops( + self, input_dict: Dict[str, Union[torch.Tensor, Any]], exclude_embeddings: bool = True + ) -> int: + """ + Get number of (optionally, non-embeddings) floating-point operations for the forward and backward passes of a + batch with this transformer model. Default approximation neglects the quadratic dependency on the number of + tokens (valid if `12 * d_model << sequence_length`) as laid out in [this + paper](https://arxiv.org/pdf/2001.08361.pdf) section 2.1. Should be overridden for transformers with parameter + re-use e.g. Albert or Universal Transformers, or if doing long-range modeling with very high sequence lengths. + + Args: + batch_size (`int`): + The batch size for the forward pass. + + sequence_length (`int`): + The number of tokens in each line of the batch. + + exclude_embeddings (`bool`, *optional*, defaults to `True`): + Whether or not to count embedding and softmax operations. + + Returns: + `int`: The number of floating-point operations. + """ + + return 6 * self.estimate_tokens(input_dict) * self.num_parameters(exclude_embeddings=exclude_embeddings) + + +class PreTrainedModel(nn.Module, ModuleUtilsMixin, GenerationMixin, PushToHubMixin): + r""" + Base class for all models. + + [`PreTrainedModel`] takes care of storing the configuration of the models and handles methods for loading, + downloading and saving models as well as a few methods common to all models to: + + - resize the input embeddings, + - prune heads in the self-attention heads. + + Class attributes (overridden by derived classes): + + - **config_class** ([`PretrainedConfig`]) -- A subclass of [`PretrainedConfig`] to use as configuration class + for this model architecture. + - **load_tf_weights** (`Callable`) -- A python *method* for loading a TensorFlow checkpoint in a PyTorch model, + taking as arguments: + + - **model** ([`PreTrainedModel`]) -- An instance of the model on which to load the TensorFlow checkpoint. + - **config** ([`PreTrainedConfig`]) -- An instance of the configuration associated to the model. + - **path** (`str`) -- A path to the TensorFlow checkpoint. + + - **base_model_prefix** (`str`) -- A string indicating the attribute associated to the base model in derived + classes of the same architecture adding modules on top of the base model. + - **is_parallelizable** (`bool`) -- A flag indicating whether this model supports model parallelization. + - **main_input_name** (`str`) -- The name of the principal input to the model (often `input_ids` for NLP + models, `pixel_values` for vision models and `input_values` for speech models). + """ + config_class = None + base_model_prefix = "" + main_input_name = "input_ids" + _auto_class = None + _no_split_modules = None + _skip_keys_device_placement = None + _keep_in_fp32_modules = None + + # a list of `re` patterns of `state_dict` keys that should be removed from the list of missing + # keys we find (keys inside the model but not in the checkpoint) and avoid unnecessary warnings. + _keys_to_ignore_on_load_missing = None + # a list of `re` patterns of `state_dict` keys that should be removed from the list of + # unexpected keys we find (keys inside the checkpoint but not the model) and avoid unnecessary + # warnings. + _keys_to_ignore_on_load_unexpected = None + # a list of `state_dict` keys to ignore when saving the model (useful for keys that aren't + # trained, but which are either deterministic or tied variables) + _keys_to_ignore_on_save = None + + is_parallelizable = False + supports_gradient_checkpointing = False + + @property + def dummy_inputs(self) -> Dict[str, torch.Tensor]: + """ + `Dict[str, torch.Tensor]`: Dummy inputs to do a forward pass in the network. + """ + return {"input_ids": torch.tensor(DUMMY_INPUTS)} + + @property + def framework(self) -> str: + """ + :str: Identifies that this is a PyTorch model. + """ + return "pt" + + def __init__(self, config: PretrainedConfig, *inputs, **kwargs): + super().__init__() + if not isinstance(config, PretrainedConfig): + raise ValueError( + f"Parameter config in `{self.__class__.__name__}(config)` should be an instance of class " + "`PretrainedConfig`. To create a model from a pretrained model use " + f"`model = {self.__class__.__name__}.from_pretrained(PRETRAINED_MODEL_NAME)`" + ) + # Save config and origin of the pretrained weights if given in model + self.config = config + self.name_or_path = config.name_or_path + self.warnings_issued = {} + self.generation_config = GenerationConfig.from_model_config(config) if self.can_generate() else None + + def post_init(self): + """ + A method executed at the end of each Transformer model initialization, to execute code that needs the model's + modules properly initialized (such as weight initialization). + """ + self.init_weights() + self._backward_compatibility_gradient_checkpointing() + + def _backward_compatibility_gradient_checkpointing(self): + if self.supports_gradient_checkpointing and getattr(self.config, "gradient_checkpointing", False): + self.gradient_checkpointing_enable() + # Remove the attribute now that is has been consumed, so it's no saved in the config. + delattr(self.config, "gradient_checkpointing") + + @classmethod + def _from_config(cls, config, **kwargs): + """ + All context managers that the model should be initialized under go here. + + Args: + torch_dtype (`torch.dtype`, *optional*): + Override the default `torch.dtype` and load the model under this dtype. + """ + torch_dtype = kwargs.pop("torch_dtype", None) + + # override default dtype if needed + dtype_orig = None + if torch_dtype is not None: + dtype_orig = cls._set_default_torch_dtype(torch_dtype) + + if is_deepspeed_zero3_enabled(): + import deepspeed + + logger.info("Detected DeepSpeed ZeRO-3: activating zero.init() for this model") + # this immediately partitions the model across all gpus, to avoid the overhead in time + # and memory copying it on CPU or each GPU first + with deepspeed.zero.Init(config_dict_or_path=deepspeed_config()): + model = cls(config, **kwargs) + else: + model = cls(config, **kwargs) + + # restore default dtype if it was modified + if dtype_orig is not None: + torch.set_default_dtype(dtype_orig) + + return model + + @classmethod + def _set_default_torch_dtype(cls, dtype: torch.dtype) -> torch.dtype: + """ + Change the default dtype and return the previous one. This is needed when wanting to instantiate the model + under specific dtype. + + Args: + dtype (`torch.dtype`): + a floating dtype to set to. + + Returns: + `torch.dtype`: the original `dtype` that can be used to restore `torch.set_default_dtype(dtype)` if it was + modified. If it wasn't, returns `None`. + + Note `set_default_dtype` currently only works with floating-point types and asserts if for example, + `torch.int64` is passed. So if a non-float `dtype` is passed this functions will throw an exception. + """ + if not dtype.is_floating_point: + raise ValueError( + f"Can't instantiate {cls.__name__} model under dtype={dtype} since it is not a floating point dtype" + ) + + logger.info(f"Instantiating {cls.__name__} model under default dtype {dtype}.") + dtype_orig = torch.get_default_dtype() + torch.set_default_dtype(dtype) + return dtype_orig + + @property + def base_model(self) -> nn.Module: + """ + `torch.nn.Module`: The main body of the model. + """ + return getattr(self, self.base_model_prefix, self) + + def can_generate(self) -> bool: + """ + Returns whether this model can generate sequences with `.generate()`. + + Returns: + `bool`: Whether this model can generate sequences with `.generate()`. + """ + # Detects whether `prepare_inputs_for_generation` has been overwritten, which is a requirement for generation + if "GenerationMixin" in str(self.prepare_inputs_for_generation.__func__): + return False + return True + + def enable_input_require_grads(self): + """ + Enables the gradients for the input embeddings. This is useful for fine-tuning adapter weights while keeping + the model weights fixed. + """ + + def make_inputs_require_grads(module, input, output): + output.requires_grad_(True) + + self._require_grads_hook = self.get_input_embeddings().register_forward_hook(make_inputs_require_grads) + + def disable_input_require_grads(self): + """ + Removes the `_require_grads_hook`. + """ + self._require_grads_hook.remove() + + def get_input_embeddings(self) -> nn.Module: + """ + Returns the model's input embeddings. + + Returns: + `nn.Module`: A torch module mapping vocabulary to hidden states. + """ + base_model = getattr(self, self.base_model_prefix, self) + if base_model is not self: + return base_model.get_input_embeddings() + else: + raise NotImplementedError + + def set_input_embeddings(self, value: nn.Module): + """ + Set model's input embeddings. + + Args: + value (`nn.Module`): A module mapping vocabulary to hidden states. + """ + base_model = getattr(self, self.base_model_prefix, self) + if base_model is not self: + base_model.set_input_embeddings(value) + else: + raise NotImplementedError + + def get_output_embeddings(self) -> nn.Module: + """ + Returns the model's output embeddings. + + Returns: + `nn.Module`: A torch module mapping hidden states to vocabulary. + """ + return None # Overwrite for models with output embeddings + + def _init_weights(self, module): + """ + Initialize the weights. This method should be overridden by derived class. + """ + pass + + def _initialize_weights(self, module): + """ + Initialize the weights if they are not already initialized. + """ + if getattr(module, "_is_hf_initialized", False): + return + self._init_weights(module) + module._is_hf_initialized = True + + def tie_weights(self): + """ + Tie the weights between the input embeddings and the output embeddings. + + If the `torchscript` flag is set in the configuration, can't handle parameter sharing so we are cloning the + weights instead. + """ + if getattr(self.config, "tie_word_embeddings", True): + output_embeddings = self.get_output_embeddings() + if output_embeddings is not None: + self._tie_or_clone_weights(output_embeddings, self.get_input_embeddings()) + + if getattr(self.config, "is_encoder_decoder", False) and getattr(self.config, "tie_encoder_decoder", False): + if hasattr(self, self.base_model_prefix): + self = getattr(self, self.base_model_prefix) + self._tie_encoder_decoder_weights(self.encoder, self.decoder, self.base_model_prefix) + + for module in self.modules(): + if hasattr(module, "_tie_weights"): + module._tie_weights() + + @staticmethod + def _tie_encoder_decoder_weights(encoder: nn.Module, decoder: nn.Module, base_model_prefix: str): + uninitialized_encoder_weights: List[str] = [] + if decoder.__class__ != encoder.__class__: + logger.info( + f"{decoder.__class__} and {encoder.__class__} are not equal. In this case make sure that all encoder" + " weights are correctly initialized." + ) + + def tie_encoder_to_decoder_recursively( + decoder_pointer: nn.Module, + encoder_pointer: nn.Module, + module_name: str, + uninitialized_encoder_weights: List[str], + depth=0, + ): + assert isinstance(decoder_pointer, nn.Module) and isinstance( + encoder_pointer, nn.Module + ), f"{decoder_pointer} and {encoder_pointer} have to be of type nn.Module" + if hasattr(decoder_pointer, "weight"): + assert hasattr(encoder_pointer, "weight") + encoder_pointer.weight = decoder_pointer.weight + if hasattr(decoder_pointer, "bias"): + assert hasattr(encoder_pointer, "bias") + encoder_pointer.bias = decoder_pointer.bias + return + + encoder_modules = encoder_pointer._modules + decoder_modules = decoder_pointer._modules + if len(decoder_modules) > 0: + assert ( + len(encoder_modules) > 0 + ), f"Encoder module {encoder_pointer} does not match decoder module {decoder_pointer}" + + all_encoder_weights = {module_name + "/" + sub_name for sub_name in encoder_modules.keys()} + encoder_layer_pos = 0 + for name, module in decoder_modules.items(): + if name.isdigit(): + encoder_name = str(int(name) + encoder_layer_pos) + decoder_name = name + if not isinstance(decoder_modules[decoder_name], type(encoder_modules[encoder_name])) and len( + encoder_modules + ) != len(decoder_modules): + # this can happen if the name corresponds to the position in a list module list of layers + # in this case the decoder has added a cross-attention that the encoder does not have + # thus skip this step and subtract one layer pos from encoder + encoder_layer_pos -= 1 + continue + elif name not in encoder_modules: + continue + elif depth > 500: + raise ValueError( + "Max depth of recursive function `tie_encoder_to_decoder` reached. It seems that there is" + " a circular dependency between two or more `nn.Modules` of your model." + ) + else: + decoder_name = encoder_name = name + tie_encoder_to_decoder_recursively( + decoder_modules[decoder_name], + encoder_modules[encoder_name], + module_name + "/" + name, + uninitialized_encoder_weights, + depth=depth + 1, + ) + all_encoder_weights.remove(module_name + "/" + encoder_name) + + uninitialized_encoder_weights += list(all_encoder_weights) + + # tie weights recursively + tie_encoder_to_decoder_recursively(decoder, encoder, base_model_prefix, uninitialized_encoder_weights) + if len(uninitialized_encoder_weights) > 0: + logger.warning( + f"The following encoder weights were not tied to the decoder {uninitialized_encoder_weights}" + ) + + def _tie_or_clone_weights(self, output_embeddings, input_embeddings): + """Tie or clone module weights depending of whether we are using TorchScript or not""" + if self.config.torchscript: + output_embeddings.weight = nn.Parameter(input_embeddings.weight.clone()) + else: + output_embeddings.weight = input_embeddings.weight + + if getattr(output_embeddings, "bias", None) is not None: + output_embeddings.bias.data = nn.functional.pad( + output_embeddings.bias.data, + ( + 0, + output_embeddings.weight.shape[0] - output_embeddings.bias.shape[0], + ), + "constant", + 0, + ) + if hasattr(output_embeddings, "out_features") and hasattr(input_embeddings, "num_embeddings"): + output_embeddings.out_features = input_embeddings.num_embeddings + + def resize_token_embeddings(self, new_num_tokens: Optional[int] = None) -> nn.Embedding: + """ + Resizes input token embeddings matrix of the model if `new_num_tokens != config.vocab_size`. + + Takes care of tying weights embeddings afterwards if the model class has a `tie_weights()` method. + + Arguments: + new_num_tokens (`int`, *optional*): + The number of new tokens in the embedding matrix. Increasing the size will add newly initialized + vectors at the end. Reducing the size will remove vectors from the end. If not provided or `None`, just + returns a pointer to the input tokens `torch.nn.Embedding` module of the model without doing anything. + + Return: + `torch.nn.Embedding`: Pointer to the input tokens Embeddings Module of the model. + """ + model_embeds = self._resize_token_embeddings(new_num_tokens) + if new_num_tokens is None: + return model_embeds + + # Update base model and current model config + self.config.vocab_size = new_num_tokens + self.vocab_size = new_num_tokens + + # Tie weights again if needed + self.tie_weights() + + return model_embeds + + def _resize_token_embeddings(self, new_num_tokens): + old_embeddings = self.get_input_embeddings() + new_embeddings = self._get_resized_embeddings(old_embeddings, new_num_tokens) + self.set_input_embeddings(new_embeddings) + + # if word embeddings are not tied, make sure that lm head is resized as well + if self.get_output_embeddings() is not None and not self.config.tie_word_embeddings: + old_lm_head = self.get_output_embeddings() + new_lm_head = self._get_resized_lm_head(old_lm_head, new_num_tokens) + self.set_output_embeddings(new_lm_head) + + return self.get_input_embeddings() + + def _get_resized_embeddings( + self, old_embeddings: nn.Embedding, new_num_tokens: Optional[int] = None + ) -> nn.Embedding: + """ + Build a resized Embedding Module from a provided token Embedding Module. Increasing the size will add newly + initialized vectors at the end. Reducing the size will remove vectors from the end + + Args: + old_embeddings (`torch.nn.Embedding`): + Old embeddings to be resized. + new_num_tokens (`int`, *optional*): + New number of tokens in the embedding matrix. + + Increasing the size will add newly initialized vectors at the end. Reducing the size will remove + vectors from the end. If not provided or `None`, just returns a pointer to the input tokens + `torch.nn.Embedding` module of the model without doing anything. + + Return: + `torch.nn.Embedding`: Pointer to the resized Embedding Module or the old Embedding Module if + `new_num_tokens` is `None` + """ + if new_num_tokens is None: + return old_embeddings + + if is_deepspeed_zero3_enabled(): + import deepspeed + + with deepspeed.zero.GatheredParameters(old_embeddings.weight, modifier_rank=None): + old_num_tokens, old_embedding_dim = old_embeddings.weight.size() + else: + old_num_tokens, old_embedding_dim = old_embeddings.weight.size() + + if old_num_tokens == new_num_tokens: + return old_embeddings + + if not isinstance(old_embeddings, nn.Embedding): + raise TypeError( + f"Old embeddings are of type {type(old_embeddings)}, which is not an instance of {nn.Embedding}. You" + " should either use a different resize function or make sure that `old_embeddings` are an instance of" + f" {nn.Embedding}." + ) + + # Build new embeddings + new_embeddings = nn.Embedding(new_num_tokens, old_embedding_dim) + new_embeddings.to(old_embeddings.weight.device, dtype=old_embeddings.weight.dtype) + + # initialize all new embeddings (in particular added tokens) + self._init_weights(new_embeddings) + + # Copy token embeddings from the previous weights + + # numbers of tokens to copy + n = min(old_num_tokens, new_num_tokens) + if is_deepspeed_zero3_enabled(): + import deepspeed + + with deepspeed.zero.GatheredParameters(old_embeddings.weight, modifier_rank=0): + if torch.distributed.get_rank() == 0: + new_embeddings.weight.data[:n, :] = old_embeddings.weight.data[:n, :] + else: + new_embeddings.weight.data[:n, :] = old_embeddings.weight.data[:n, :] + + return new_embeddings + + def _get_resized_lm_head( + self, old_lm_head: nn.Linear, new_num_tokens: Optional[int] = None, transposed: Optional[bool] = False + ) -> nn.Linear: + """ + Build a resized Linear Module from a provided old Linear Module. Increasing the size will add newly initialized + vectors at the end. Reducing the size will remove vectors from the end + + Args: + old_lm_head (`torch.nn.Linear`): + Old lm head liner layer to be resized. + new_num_tokens (`int`, *optional*): + New number of tokens in the linear matrix. + + Increasing the size will add newly initialized vectors at the end. Reducing the size will remove + vectors from the end. If not provided or `None`, just returns a pointer to the input tokens + `torch.nn.Linear` module of the model without doing anything. transposed (`bool`, *optional*, defaults + to `False`): Whether `old_lm_head` is transposed or not. If True `old_lm_head.size()` is `lm_head_dim, + vocab_size` else `vocab_size, lm_head_dim`. + + Return: + `torch.nn.Linear`: Pointer to the resized Linear Module or the old Linear Module if `new_num_tokens` is + `None` + """ + if new_num_tokens is None: + return old_lm_head + + if is_deepspeed_zero3_enabled(): + import deepspeed + + with deepspeed.zero.GatheredParameters(old_lm_head.weight, modifier_rank=None): + old_num_tokens, old_lm_head_dim = ( + old_lm_head.weight.size() if not transposed else old_lm_head.weight.t().size() + ) + else: + old_num_tokens, old_lm_head_dim = ( + old_lm_head.weight.size() if not transposed else old_lm_head.weight.t().size() + ) + + if old_num_tokens == new_num_tokens: + return old_lm_head + + if not isinstance(old_lm_head, nn.Linear): + raise TypeError( + f"Old language model head is of type {type(old_lm_head)}, which is not an instance of {nn.Linear}. You" + " should either use a different resize function or make sure that `old_lm_head` are an instance of" + f" {nn.Linear}." + ) + + # Build new lm head + new_lm_head_shape = (old_lm_head_dim, new_num_tokens) if not transposed else (new_num_tokens, old_lm_head_dim) + has_new_lm_head_bias = old_lm_head.bias is not None + new_lm_head = nn.Linear(*new_lm_head_shape, bias=has_new_lm_head_bias) + new_lm_head = new_lm_head.to(old_lm_head.weight.device, dtype=old_lm_head.weight.dtype) + + # initialize new lm head (in particular added tokens) + self._init_weights(new_lm_head) + + num_tokens_to_copy = min(old_num_tokens, new_num_tokens) + + # XXX: put the long block of code in a wrapper + if is_deepspeed_zero3_enabled(): + import deepspeed + + params = [old_lm_head.weight, old_lm_head.bias, new_lm_head.weight, new_lm_head.bias] + with deepspeed.zero.GatheredParameters(params, modifier_rank=0): + if torch.distributed.get_rank() == 0: + # Copy old lm head weights to new lm head + if not transposed: + new_lm_head.weight.data[:num_tokens_to_copy, :] = old_lm_head.weight.data[ + :num_tokens_to_copy, : + ] + else: + new_lm_head.weight.data[:, :num_tokens_to_copy] = old_lm_head.weight.data[ + :, :num_tokens_to_copy + ] + + # Copy bias weights to new lm head + if has_new_lm_head_bias: + new_lm_head.bias.data[:num_tokens_to_copy] = old_lm_head.bias.data[:num_tokens_to_copy] + else: + # Copy old lm head weights to new lm head + if not transposed: + new_lm_head.weight.data[:num_tokens_to_copy, :] = old_lm_head.weight.data[:num_tokens_to_copy, :] + else: + new_lm_head.weight.data[:, :num_tokens_to_copy] = old_lm_head.weight.data[:, :num_tokens_to_copy] + + # Copy bias weights to new lm head + if has_new_lm_head_bias: + new_lm_head.bias.data[:num_tokens_to_copy] = old_lm_head.bias.data[:num_tokens_to_copy] + + return new_lm_head + + def resize_position_embeddings(self, new_num_position_embeddings: int): + raise NotImplementedError( + f"`resize_position_embeddings` is not implemented for {self.__class__}`. To implement it, you should " + f"overwrite this method in the class {self.__class__} in `modeling_{self.__class__.__module__}.py`" + ) + + def get_position_embeddings(self) -> Union[nn.Embedding, Tuple[nn.Embedding]]: + raise NotImplementedError( + f"`get_position_embeddings` is not implemented for {self.__class__}`. To implement it, you should " + f"overwrite this method in the class {self.__class__} in `modeling_{self.__class__.__module__}.py`" + ) + + def init_weights(self): + """ + If needed prunes and maybe initializes weights. If using a custom `PreTrainedModel`, you need to implement any + initialization logic in `_init_weights`. + """ + # Prune heads if needed + if self.config.pruned_heads: + self.prune_heads(self.config.pruned_heads) + + if _init_weights: + # Initialize weights + self.apply(self._initialize_weights) + + # Tie weights should be skipped when not initializing all weights + # since from_pretrained(...) calls tie weights anyways + self.tie_weights() + + def prune_heads(self, heads_to_prune: Dict[int, List[int]]): + """ + Prunes heads of the base model. + + Arguments: + heads_to_prune (`Dict[int, List[int]]`): + Dictionary with keys being selected layer indices (`int`) and associated values being the list of heads + to prune in said layer (list of `int`). For instance {1: [0, 2], 2: [2, 3]} will prune heads 0 and 2 on + layer 1 and heads 2 and 3 on layer 2. + """ + # save new sets of pruned heads as union of previously stored pruned heads and newly pruned heads + for layer, heads in heads_to_prune.items(): + union_heads = set(self.config.pruned_heads.get(layer, [])) | set(heads) + self.config.pruned_heads[layer] = list(union_heads) # Unfortunately we have to store it as list for JSON + + self.base_model._prune_heads(heads_to_prune) + + def gradient_checkpointing_enable(self): + """ + Activates gradient checkpointing for the current model. + + Note that in other frameworks this feature can be referred to as "activation checkpointing" or "checkpoint + activations". + """ + if not self.supports_gradient_checkpointing: + raise ValueError(f"{self.__class__.__name__} does not support gradient checkpointing.") + self.apply(partial(self._set_gradient_checkpointing, value=True)) + + def gradient_checkpointing_disable(self): + """ + Deactivates gradient checkpointing for the current model. + + Note that in other frameworks this feature can be referred to as "activation checkpointing" or "checkpoint + activations". + """ + if self.supports_gradient_checkpointing: + self.apply(partial(self._set_gradient_checkpointing, value=False)) + + @property + def is_gradient_checkpointing(self) -> bool: + """ + Whether gradient checkpointing is activated for this model or not. + + Note that in other frameworks this feature can be referred to as "activation checkpointing" or "checkpoint + activations". + """ + return any(hasattr(m, "gradient_checkpointing") and m.gradient_checkpointing for m in self.modules()) + + def save_pretrained( + self, + save_directory: Union[str, os.PathLike], + is_main_process: bool = True, + state_dict: Optional[dict] = None, + save_function: Callable = torch.save, + push_to_hub: bool = False, + max_shard_size: Union[int, str] = "10GB", + safe_serialization: bool = False, + variant: Optional[str] = None, + **kwargs, + ): + """ + Save a model and its configuration file to a directory, so that it can be re-loaded using the + [`~PreTrainedModel.from_pretrained`] class method. + + Arguments: + save_directory (`str` or `os.PathLike`): + Directory to which to save. Will be created if it doesn't exist. + is_main_process (`bool`, *optional*, defaults to `True`): + Whether the process calling this is the main process or not. Useful when in distributed training like + TPUs and need to call this function on all processes. In this case, set `is_main_process=True` only on + the main process to avoid race conditions. + state_dict (nested dictionary of `torch.Tensor`): + The state dictionary of the model to save. Will default to `self.state_dict()`, but can be used to only + save parts of the model or if special precautions need to be taken when recovering the state dictionary + of a model (like when using model parallelism). + save_function (`Callable`): + The function to use to save the state dictionary. Useful on distributed training like TPUs when one + need to replace `torch.save` by another method. + push_to_hub (`bool`, *optional*, defaults to `False`): + Whether or not to push your model to the Hugging Face model hub after saving it. You can specify the + repository you want to push to with `repo_id` (will default to the name of `save_directory` in your + namespace). + max_shard_size (`int` or `str`, *optional*, defaults to `"10GB"`): + The maximum size for a checkpoint before being sharded. Checkpoints shard will then be each of size + lower than this size. If expressed as a string, needs to be digits followed by a unit (like `"5MB"`). + + + + If a single weight of the model is bigger than `max_shard_size`, it will be in its own checkpoint shard + which will be bigger than `max_shard_size`. + + + + safe_serialization (`bool`, *optional*, defaults to `False`): + Whether to save the model using `safetensors` or the traditional PyTorch way (that uses `pickle`). + variant (`str`, *optional*): + If specified, weights are saved in the format pytorch_model..bin. + + kwargs: + Additional key word arguments passed along to the [`~utils.PushToHubMixin.push_to_hub`] method. + """ + # Checks if the model has been loaded in 8-bit + if getattr(self, "is_loaded_in_8bit", False) and getattr(self, "is_8bit_serializable", False): + warnings.warn( + "You are calling `save_pretrained` to a 8-bit converted model you may likely encounter unexepected" + " behaviors. If you want to save 8-bit models, make sure to have `bitsandbytes>0.37.2` installed.", + UserWarning, + ) + + if getattr(self, "is_loaded_in_4bit", False): + raise NotImplementedError( + "You are calling `save_pretrained` on a 4-bit converted model. This is currently not supported" + ) + + if "save_config" in kwargs: + warnings.warn( + "`save_config` is deprecated and will be removed in v5 of Transformers. Use `is_main_process` instead." + ) + is_main_process = kwargs.pop("save_config") + if safe_serialization and not is_safetensors_available(): + raise ImportError("`safe_serialization` requires the `safetensors library: `pip install safetensors`.") + + if os.path.isfile(save_directory): + logger.error(f"Provided path ({save_directory}) should be a directory, not a file") + return + + os.makedirs(save_directory, exist_ok=True) + + if push_to_hub: + commit_message = kwargs.pop("commit_message", None) + repo_id = kwargs.pop("repo_id", save_directory.split(os.path.sep)[-1]) + repo_id = self._create_repo(repo_id, **kwargs) + files_timestamps = self._get_files_timestamps(save_directory) + + # Only save the model itself if we are using distributed training + model_to_save = unwrap_model(self) + + # save the string version of dtype to the config, e.g. convert torch.float32 => "float32" + # we currently don't use this setting automatically, but may start to use with v5 + dtype = get_parameter_dtype(model_to_save) + model_to_save.config.torch_dtype = str(dtype).split(".")[1] + + # Attach architecture to the config + model_to_save.config.architectures = [model_to_save.__class__.__name__] + + # If we have a custom model, we copy the file defining it in the folder and set the attributes so it can be + # loaded from the Hub. + if self._auto_class is not None: + custom_object_save(self, save_directory, config=self.config) + + # Save the config + if is_main_process: + model_to_save.config.save_pretrained(save_directory) + if self.can_generate(): + model_to_save.generation_config.save_pretrained(save_directory) + + # Save the model + if state_dict is None: + state_dict = model_to_save.state_dict() + + # Translate state_dict from smp to hf if saving with smp >= 1.10 + if IS_SAGEMAKER_MP_POST_1_10: + for smp_to_hf, _ in smp.state.module_manager.translate_functions: + state_dict = smp_to_hf(state_dict) + + # Handle the case where some state_dict keys shouldn't be saved + if self._keys_to_ignore_on_save is not None: + for ignore_key in self._keys_to_ignore_on_save: + if ignore_key in state_dict.keys(): + del state_dict[ignore_key] + if safe_serialization: + # Safetensors does not allow tensor aliasing. + # We're going to remove aliases before saving + ptrs = collections.defaultdict(list) + for name, tensor in state_dict.items(): + ident = (tensor.data_ptr(), tensor.device, tensor.shape, tensor.stride()) + ptrs[ident].append(name) + + # These are all the pointers of shared tensors. + shared_ptrs = {ptr: names for ptr, names in ptrs.items() if len(names) > 1} + warn_names = set() + for names in shared_ptrs.values(): + # Removing the keys which are declared as known duplicates on + # load. This allows to make sure the name which is kept is consistent. + if self._keys_to_ignore_on_load_missing is not None: + found = 0 + for name in sorted(names): + matches_pattern = any(re.search(pat, name) for pat in self._keys_to_ignore_on_load_missing) + if matches_pattern and name in state_dict: + found += 1 + if found < len(names): + del state_dict[name] + + # When not all duplicates have been cleaned, still remove those keys, but put a clear warning. + # If the link between tensors was done at runtime then `from_pretrained` will not get + # the key back leading to random tensor. A proper warning will be shown + # during reload (if applicable), but since the file is not necessarily compatible with + # the config, better show a proper warning. + found = 0 + for name in names: + if name in state_dict: + found += 1 + if found > 1: + del state_dict[name] + warn_names.add(name) + if len(warn_names) > 0: + logger.warning_once( + f"Removed shared tensor {warn_names} while saving. This should be OK, but check by verifying that you don't receive any warning while reloading", + ) + + # Shard the model if it is too big. + weights_name = SAFE_WEIGHTS_NAME if safe_serialization else WEIGHTS_NAME + weights_name = _add_variant(weights_name, variant) + + shards, index = shard_checkpoint(state_dict, max_shard_size=max_shard_size, weights_name=weights_name) + + # Clean the folder from a previous save + for filename in os.listdir(save_directory): + full_filename = os.path.join(save_directory, filename) + # If we have a shard file that is not going to be replaced, we delete it, but only from the main process + # in distributed settings to avoid race conditions. + weights_no_suffix = weights_name.replace(".bin", "").replace(".safetensors", "") + + # make sure that file to be deleted matches format of sharded file, e.g. pytorch_model-00001-of-00005 + filename_no_suffix = filename.replace(".bin", "").replace(".safetensors", "") + reg = re.compile(r"(.*?)-\d{5}-of-\d{5}") + + if ( + filename.startswith(weights_no_suffix) + and os.path.isfile(full_filename) + and filename not in shards.keys() + and is_main_process + and reg.fullmatch(filename_no_suffix) is not None + ): + os.remove(full_filename) + + # Save the model + for shard_file, shard in shards.items(): + if safe_serialization: + # At some point we will need to deal better with save_function (used for TPU and other distributed + # joyfulness), but for now this enough. + safe_save_file(shard, os.path.join(save_directory, shard_file), metadata={"format": "pt"}) + else: + save_function(shard, os.path.join(save_directory, shard_file)) + + if index is None: + path_to_weights = os.path.join(save_directory, _add_variant(WEIGHTS_NAME, variant)) + logger.info(f"Model weights saved in {path_to_weights}") + else: + save_index_file = SAFE_WEIGHTS_INDEX_NAME if safe_serialization else WEIGHTS_INDEX_NAME + save_index_file = os.path.join(save_directory, _add_variant(save_index_file, variant)) + # Save the index as well + with open(save_index_file, "w", encoding="utf-8") as f: + content = json.dumps(index, indent=2, sort_keys=True) + "\n" + f.write(content) + logger.info( + f"The model is bigger than the maximum size per checkpoint ({max_shard_size}) and is going to be " + f"split in {len(shards)} checkpoint shards. You can find where each parameters has been saved in the " + f"index located at {save_index_file}." + ) + + if push_to_hub: + self._upload_modified_files( + save_directory, + repo_id, + files_timestamps, + commit_message=commit_message, + token=kwargs.get("use_auth_token"), + ) + + def get_memory_footprint(self, return_buffers=True): + r""" + Get the memory footprint of a model. This will return the memory footprint of the current model in bytes. + Useful to benchmark the memory footprint of the current model and design some tests. Solution inspired from the + PyTorch discussions: https://discuss.pytorch.org/t/gpu-memory-that-model-uses/56822/2 + + Arguments: + return_buffers (`bool`, *optional*, defaults to `True`): + Whether to return the size of the buffer tensors in the computation of the memory footprint. Buffers + are tensors that do not require gradients and not registered as parameters. E.g. mean and std in batch + norm layers. Please see: https://discuss.pytorch.org/t/what-pytorch-means-by-buffers/120266/2 + """ + mem = sum([param.nelement() * param.element_size() for param in self.parameters()]) + if return_buffers: + mem_bufs = sum([buf.nelement() * buf.element_size() for buf in self.buffers()]) + mem = mem + mem_bufs + return mem + + def to(self, *args, **kwargs): + # Checks if the model has been loaded in 8-bit + if getattr(self, "is_quantized", False): + raise ValueError( + "`.to` is not supported for `4-bit` or `8-bit` models. Please use the model as it is, since the" + " model has already been set to the correct devices and casted to the correct `dtype`." + ) + else: + return super().to(*args, **kwargs) + + def half(self, *args): + # Checks if the model has been loaded in 8-bit + if getattr(self, "is_quantized", False): + raise ValueError( + "`.half()` is not supported for `4-bit` or `8-bit` models. Please use the model as it is, since the" + " model has already been casted to the correct `dtype`." + ) + else: + return super().half(*args) + + def float(self, *args): + # Checks if the model has been loaded in 8-bit + if getattr(self, "is_quantized", False): + raise ValueError( + "`.float()` is not supported for `4-bit` or `8-bit` models. Please use the model as it is, since the" + " model has already been casted to the correct `dtype`." + ) + else: + return super().float(*args) + + @classmethod + def from_pretrained(cls, pretrained_model_name_or_path: Optional[Union[str, os.PathLike]], *model_args, **kwargs): + r""" + Instantiate a pretrained pytorch model from a pre-trained model configuration. + + The model is set in evaluation mode by default using `model.eval()` (Dropout modules are deactivated). To train + the model, you should first set it back in training mode with `model.train()`. + + The warning *Weights from XXX not initialized from pretrained model* means that the weights of XXX do not come + pretrained with the rest of the model. It is up to you to train those weights with a downstream fine-tuning + task. + + The warning *Weights from XXX not used in YYY* means that the layer XXX is not used by YYY, therefore those + weights are discarded. + + Parameters: + pretrained_model_name_or_path (`str` or `os.PathLike`, *optional*): + Can be either: + + - A string, the *model id* of a pretrained model hosted inside a model repo on huggingface.co. + Valid model ids can be located at the root-level, like `bert-base-uncased`, or namespaced under a + user or organization name, like `dbmdz/bert-base-german-cased`. + - A path to a *directory* containing model weights saved using + [`~PreTrainedModel.save_pretrained`], e.g., `./my_model_directory/`. + - A path or url to a *tensorflow index checkpoint file* (e.g, `./tf_model/model.ckpt.index`). In + this case, `from_tf` should be set to `True` and a configuration object should be provided as + `config` argument. This loading path is slower than converting the TensorFlow checkpoint in a + PyTorch model using the provided conversion scripts and loading the PyTorch model afterwards. + - A path or url to a model folder containing a *flax checkpoint file* in *.msgpack* format (e.g, + `./flax_model/` containing `flax_model.msgpack`). In this case, `from_flax` should be set to + `True`. + - `None` if you are both providing the configuration and state dictionary (resp. with keyword + arguments `config` and `state_dict`). + model_args (sequence of positional arguments, *optional*): + All remaining positional arguments will be passed to the underlying model's `__init__` method. + config (`Union[PretrainedConfig, str, os.PathLike]`, *optional*): + Can be either: + + - an instance of a class derived from [`PretrainedConfig`], + - a string or path valid as input to [`~PretrainedConfig.from_pretrained`]. + + Configuration for the model to use instead of an automatically loaded configuration. Configuration can + be automatically loaded when: + + - The model is a model provided by the library (loaded with the *model id* string of a pretrained + model). + - The model was saved using [`~PreTrainedModel.save_pretrained`] and is reloaded by supplying the + save directory. + - The model is loaded by supplying a local directory as `pretrained_model_name_or_path` and a + configuration JSON file named *config.json* is found in the directory. + state_dict (`Dict[str, torch.Tensor]`, *optional*): + A state dictionary to use instead of a state dictionary loaded from saved weights file. + + This option can be used if you want to create a model from a pretrained configuration but load your own + weights. In this case though, you should check if using [`~PreTrainedModel.save_pretrained`] and + [`~PreTrainedModel.from_pretrained`] is not a simpler option. + cache_dir (`Union[str, os.PathLike]`, *optional*): + Path to a directory in which a downloaded pretrained model configuration should be cached if the + standard cache should not be used. + from_tf (`bool`, *optional*, defaults to `False`): + Load the model weights from a TensorFlow checkpoint save file (see docstring of + `pretrained_model_name_or_path` argument). + from_flax (`bool`, *optional*, defaults to `False`): + Load the model weights from a Flax checkpoint save file (see docstring of + `pretrained_model_name_or_path` argument). + ignore_mismatched_sizes (`bool`, *optional*, defaults to `False`): + Whether or not to raise an error if some of the weights from the checkpoint do not have the same size + as the weights of the model (if for instance, you are instantiating a model with 10 labels from a + checkpoint with 3 labels). + force_download (`bool`, *optional*, defaults to `False`): + Whether or not to force the (re-)download of the model weights and configuration files, overriding the + cached versions if they exist. + resume_download (`bool`, *optional*, defaults to `False`): + Whether or not to delete incompletely received files. Will attempt to resume the download if such a + file exists. + proxies (`Dict[str, str]`, *optional*): + A dictionary of proxy servers to use by protocol or endpoint, e.g., `{'http': 'foo.bar:3128', + 'http://hostname': 'foo.bar:4012'}`. The proxies are used on each request. + output_loading_info(`bool`, *optional*, defaults to `False`): + Whether ot not to also return a dictionary containing missing keys, unexpected keys and error messages. + local_files_only(`bool`, *optional*, defaults to `False`): + Whether or not to only look at local files (i.e., do not try to download the model). + use_auth_token (`str` or `bool`, *optional*): + The token to use as HTTP bearer authorization for remote files. If `True`, or not specified, will use + the token generated when running `huggingface-cli login` (stored in `~/.huggingface`). + revision (`str`, *optional*, defaults to `"main"`): + The specific model version to use. It can be a branch name, a tag name, or a commit id, since we use a + git-based system for storing models and other artifacts on huggingface.co, so `revision` can be any + identifier allowed by git. + + + + To test a pull request you made on the Hub, you can pass `revision="refs/pr/". + + + + mirror (`str`, *optional*): + Mirror source to accelerate downloads in China. If you are from China and have an accessibility + problem, you can set this option to resolve it. Note that we do not guarantee the timeliness or safety. + Please refer to the mirror site for more information. + _fast_init(`bool`, *optional*, defaults to `True`): + Whether or not to disable fast initialization. + + + + One should only disable *_fast_init* to ensure backwards compatibility with `transformers.__version__ < + 4.6.0` for seeded model initialization. This argument will be removed at the next major version. See + [pull request 11471](https://github.com/huggingface/transformers/pull/11471) for more information. + + + + > Parameters for big model inference + + low_cpu_mem_usage(`bool`, *optional*): + Tries to not use more than 1x model size in CPU memory (including peak memory) while loading the model. + This is an experimental feature and a subject to change at any moment. + torch_dtype (`str` or `torch.dtype`, *optional*): + Override the default `torch.dtype` and load the model under a specific `dtype`. The different options + are: + + 1. `torch.float16` or `torch.bfloat16` or `torch.float`: load in a specified + `dtype`, ignoring the model's `config.torch_dtype` if one exists. If not specified + - the model will get loaded in `torch.float` (fp32). + + 2. `"auto"` - A `torch_dtype` entry in the `config.json` file of the model will be + attempted to be used. If this entry isn't found then next check the `dtype` of the first weight in + the checkpoint that's of a floating point type and use that as `dtype`. This will load the model + using the `dtype` it was saved in at the end of the training. It can't be used as an indicator of how + the model was trained. Since it could be trained in one of half precision dtypes, but saved in fp32. + + + + For some models the `dtype` they were trained in is unknown - you may try to check the model's paper or + reach out to the authors and ask them to add this information to the model's card and to insert the + `torch_dtype` entry in `config.json` on the hub. + + + + device_map (`str` or `Dict[str, Union[int, str, torch.device]]` or `int` or `torch.device`, *optional*): + A map that specifies where each submodule should go. It doesn't need to be refined to each + parameter/buffer name, once a given module name is inside, every submodule of it will be sent to the + same device. If we only pass the device (*e.g.*, `"cpu"`, `"cuda:1"`, `"mps"`, or a GPU ordinal rank + like `1`) on which the model will be allocated, the device map will map the entire model to this + device. Passing `device_map = 0` means put the whole model on GPU 0. + + To have Accelerate compute the most optimized `device_map` automatically, set `device_map="auto"`. For + more information about each option see [designing a device + map](https://hf.co/docs/accelerate/main/en/usage_guides/big_modeling#designing-a-device-map). + max_memory (`Dict`, *optional*): + A dictionary device identifier to maximum memory. Will default to the maximum memory available for each + GPU and the available CPU RAM if unset. + offload_folder (`str` or `os.PathLike`, *optional*): + If the `device_map` contains any value `"disk"`, the folder where we will offload weights. + offload_state_dict (`bool`, *optional*): + If `True`, will temporarily offload the CPU state dict to the hard drive to avoid getting out of CPU + RAM if the weight of the CPU state dict + the biggest shard of the checkpoint does not fit. Defaults to + `True` when there is some disk offload. + load_in_8bit (`bool`, *optional*, defaults to `False`): + If `True`, will convert the loaded model into mixed-8bit quantized model. To use this feature please + install `bitsandbytes` compiled with your CUDA version by running `pip install -i + https://test.pypi.org/simple/ bitsandbytes-cudaXXX` where XXX is your CUDA version (e.g. 11.6 = 116). + Make also sure that you have enough GPU RAM to store half of the model size since the 8bit modules are + not compiled and adapted for CPUs. + quantization_config (`Dict`, *optional*): + A dictionary of configuration parameters for the `bitsandbytes` library and loading the model using + advanced features such as offloading in fp32 on CPU or on disk. + subfolder (`str`, *optional*, defaults to `""`): + In case the relevant files are located inside a subfolder of the model repo on huggingface.co, you can + specify the folder name here. + variant (`str`, *optional*): + If specified load weights from `variant` filename, *e.g.* pytorch_model..bin. `variant` is + ignored when using `from_tf` or `from_flax`. + + kwargs (remaining dictionary of keyword arguments, *optional*): + Can be used to update the configuration object (after it being loaded) and initiate the model (e.g., + `output_attentions=True`). Behaves differently depending on whether a `config` is provided or + automatically loaded: + + - If a configuration is provided with `config`, `**kwargs` will be directly passed to the + underlying model's `__init__` method (we assume all relevant updates to the configuration have + already been done) + - If a configuration is not provided, `kwargs` will be first passed to the configuration class + initialization function ([`~PretrainedConfig.from_pretrained`]). Each key of `kwargs` that + corresponds to a configuration attribute will be used to override said attribute with the + supplied `kwargs` value. Remaining keys that do not correspond to any configuration attribute + will be passed to the underlying model's `__init__` function. + + + + Activate the special ["offline-mode"](https://huggingface.co/transformers/installation.html#offline-mode) to + use this method in a firewalled environment. + + + + Examples: + + ```python + >>> from transformers import BertConfig, BertModel + + >>> # Download model and configuration from huggingface.co and cache. + >>> model = BertModel.from_pretrained("bert-base-uncased") + >>> # Model was saved using *save_pretrained('./test/saved_model/')* (for example purposes, not runnable). + >>> model = BertModel.from_pretrained("./test/saved_model/") + >>> # Update configuration during loading. + >>> model = BertModel.from_pretrained("bert-base-uncased", output_attentions=True) + >>> assert model.config.output_attentions == True + >>> # Loading from a TF checkpoint file instead of a PyTorch model (slower, for example purposes, not runnable). + >>> config = BertConfig.from_json_file("./tf_model/my_tf_model_config.json") + >>> model = BertModel.from_pretrained("./tf_model/my_tf_checkpoint.ckpt.index", from_tf=True, config=config) + >>> # Loading from a Flax checkpoint file instead of a PyTorch model (slower) + >>> model = BertModel.from_pretrained("bert-base-uncased", from_flax=True) + ``` + + * `low_cpu_mem_usage` algorithm: + + This is an experimental function that loads the model using ~1x model size CPU memory + + Here is how it works: + + 1. save which state_dict keys we have + 2. drop state_dict before the model is created, since the latter takes 1x model size CPU memory + 3. after the model has been instantiated switch to the meta device all params/buffers that + are going to be replaced from the loaded state_dict + 4. load state_dict 2nd time + 5. replace the params/buffers from the state_dict + + Currently, it can't handle deepspeed ZeRO stage 3 and ignores loading errors + + """ + config = kwargs.pop("config", None) + state_dict = kwargs.pop("state_dict", None) + cache_dir = kwargs.pop("cache_dir", None) + from_tf = kwargs.pop("from_tf", False) + from_flax = kwargs.pop("from_flax", False) + ignore_mismatched_sizes = kwargs.pop("ignore_mismatched_sizes", False) + force_download = kwargs.pop("force_download", False) + resume_download = kwargs.pop("resume_download", False) + proxies = kwargs.pop("proxies", None) + output_loading_info = kwargs.pop("output_loading_info", False) + local_files_only = kwargs.pop("local_files_only", False) + use_auth_token = kwargs.pop("use_auth_token", None) + revision = kwargs.pop("revision", None) + trust_remote_code = kwargs.pop("trust_remote_code", None) + _ = kwargs.pop("mirror", None) + from_pipeline = kwargs.pop("_from_pipeline", None) + from_auto_class = kwargs.pop("_from_auto", False) + _fast_init = kwargs.pop("_fast_init", True) + torch_dtype = kwargs.pop("torch_dtype", None) + low_cpu_mem_usage = kwargs.pop("low_cpu_mem_usage", None) + device_map = kwargs.pop("device_map", None) + max_memory = kwargs.pop("max_memory", None) + offload_folder = kwargs.pop("offload_folder", None) + offload_state_dict = kwargs.pop("offload_state_dict", False) + load_in_8bit = kwargs.pop("load_in_8bit", False) + load_in_4bit = kwargs.pop("load_in_4bit", False) + quantization_config = kwargs.pop("quantization_config", None) + subfolder = kwargs.pop("subfolder", "") + commit_hash = kwargs.pop("_commit_hash", None) + variant = kwargs.pop("variant", None) + use_safetensors = kwargs.pop("use_safetensors", None if is_safetensors_available() else False) + + if is_bitsandbytes_available(): + is_8bit_serializable = version.parse(importlib_metadata.version("bitsandbytes")) > version.parse("0.37.2") + else: + is_8bit_serializable = False + + if trust_remote_code is True: + logger.warning( + "The argument `trust_remote_code` is to be used with Auto classes. It has no effect here and is" + " ignored." + ) + + # change device_map into a map if we passed an int, a str or a torch.device + if isinstance(device_map, torch.device): + device_map = {"": device_map} + elif isinstance(device_map, str) and device_map not in ["auto", "balanced", "balanced_low_0", "sequential"]: + try: + device_map = {"": torch.device(device_map)} + except RuntimeError: + raise ValueError( + "When passing device_map as a string, the value needs to be a device name (e.g. cpu, cuda:0) or " + f"'auto', 'balanced', 'balanced_low_0', 'sequential' but found {device_map}." + ) + elif isinstance(device_map, int): + if device_map < 0: + raise ValueError( + "You can't pass device_map as a negative int. If you want to put the model on the cpu, pass device_map = 'cpu' " + ) + else: + device_map = {"": device_map} + + if device_map is not None: + if low_cpu_mem_usage is None: + low_cpu_mem_usage = True + elif not low_cpu_mem_usage: + raise ValueError("Passing along a `device_map` requires `low_cpu_mem_usage=True`") + + if low_cpu_mem_usage: + if device_map is not None: + # The max memory utils require PyTorch >= 1.10 to have torch.cuda.mem_get_info. + require_version_core("torch>=1.10") + + if is_deepspeed_zero3_enabled(): + raise ValueError( + "DeepSpeed Zero-3 is not compatible with `low_cpu_mem_usage=True` or with passing a `device_map`." + ) + elif not is_accelerate_available(): + raise ImportError( + "Using `low_cpu_mem_usage=True` or a `device_map` requires Accelerate: `pip install accelerate`" + ) + + if quantization_config is None: + quantization_config, kwargs = BitsAndBytesConfig.from_dict( + config_dict={"load_in_8bit": load_in_8bit, "load_in_4bit": load_in_4bit}, + return_unused_kwargs=True, + **kwargs, + ) + elif quantization_config is not None: + load_in_8bit = quantization_config.load_in_8bit + load_in_4bit = quantization_config.load_in_4bit + + quantization_config_kwargs = { + k: v for k, v in kwargs.items() if k in inspect.signature(BitsAndBytesConfig).parameters + } + + if len(quantization_config_kwargs) > 0: + raise ValueError( + "You can't pass `load_in_8bit` or any other `BitsAndBytesConfig` argument as a kwarg when passing " + "`quantization_config` argument at the same time." + ) + + if load_in_8bit or load_in_4bit: + if not (is_accelerate_available() and is_bitsandbytes_available()): + raise ImportError( + "Using `load_in_8bit=True` requires Accelerate: `pip install accelerate` and the latest version of" + " bitsandbytes `pip install -i https://test.pypi.org/simple/ bitsandbytes` or" + " pip install bitsandbytes` " + ) + + if torch_dtype is None: + # We force the `dtype` to be float16, this is a requirement from `bitsandbytes` + logger.info( + f"Overriding torch_dtype={torch_dtype} with `torch_dtype=torch.float16` due to " + "requirements of `bitsandbytes` to enable model loading in 8-bit or 4-bit. " + "Pass your own torch_dtype to specify the dtype of the remaining non-linear layers or pass" + " torch_dtype=torch.float16 to remove this warning." + ) + torch_dtype = torch.float16 + + if device_map is None: + if torch.cuda.is_available(): + device_map = {"": torch.cuda.current_device()} + else: + raise RuntimeError("No GPU found. A GPU is needed for quantization.") + logger.info( + "The device_map was not initialized." + "Setting device_map to {'':torch.cuda.current_device()}." + "If you want to use the model for inference, please set device_map ='auto' " + ) + if low_cpu_mem_usage is None: + low_cpu_mem_usage = True + + if from_tf or from_flax: + raise ValueError( + "Converting into 4-bit or 8-bit weights from tf/flax weights is currently not supported, please make" + " sure the weights are in PyTorch format." + ) + + from_pt = not (from_tf | from_flax) + + user_agent = {"file_type": "model", "framework": "pytorch", "from_auto_class": from_auto_class} + if from_pipeline is not None: + user_agent["using_pipeline"] = from_pipeline + + if is_offline_mode() and not local_files_only: + logger.info("Offline mode: forcing local_files_only=True") + local_files_only = True + + # Load config if we don't provide a configuration + if not isinstance(config, PretrainedConfig): + config_path = config if config is not None else pretrained_model_name_or_path + config, model_kwargs = cls.config_class.from_pretrained( + config_path, + cache_dir=cache_dir, + return_unused_kwargs=True, + force_download=force_download, + resume_download=resume_download, + proxies=proxies, + local_files_only=local_files_only, + use_auth_token=use_auth_token, + revision=revision, + subfolder=subfolder, + _from_auto=from_auto_class, + _from_pipeline=from_pipeline, + **kwargs, + ) + else: + model_kwargs = kwargs + + if is_8bit_serializable and quantization_config is not None and load_in_8bit: + if hasattr(config, "quantization_config"): + logger.warning( + "You passed `quantization_config` to `from_pretrained` but the model you're loading already has a" + " `quantization_config` attribute. The `quantization_config` attribute will be overwritten with the" + " one you passed to `from_pretrained`." + ) + config.quantization_config = quantization_config + elif is_8bit_serializable and not load_in_8bit and hasattr(config, "quantization_config"): + quantization_config = config.quantization_config + if isinstance(quantization_config, dict): + quantization_config = BitsAndBytesConfig.from_dict(quantization_config, return_unused_kwargs=False) + elif isinstance(quantization_config, BitsAndBytesConfig): + pass + else: + raise ValueError( + f"Invalid type for `quantization_config`: {type(quantization_config)}. Should be a `dict` or a" + " `BitsAndBytesConfig` instance." + ) + + load_in_8bit = quantization_config.load_in_8bit + + if load_in_8bit: + if torch_dtype is None: + torch_dtype = torch.float16 + if device_map is None: + if torch.cuda.is_available(): + device_map = {"": torch.cuda.current_device()} + else: + raise RuntimeError("No GPU found. A GPU is needed for quantization.") + logger.info( + "The device_map was not initialized." + "Setting device_map to {'':torch.cuda.current_device()}." + "If you want to use the model for inference, please set device_map ='auto' " + ) + if low_cpu_mem_usage is None: + low_cpu_mem_usage = True + + elif not is_8bit_serializable and not load_in_8bit and hasattr(config, "quantization_config"): + logger.warning( + "Detected the presence of a `quantization_config` attribute in the model's configuration but you don't have the correct" + " `bitsandbytes` version to support int8 serialization. Please install the latest version of `bitsandbytes` with " + " `pip install --upgrade bitsandbytes`." + ) + + if commit_hash is None: + commit_hash = getattr(config, "_commit_hash", None) + + # This variable will flag if we're loading a sharded checkpoint. In this case the archive file is just the + # index of the files. + is_sharded = False + sharded_metadata = None + # Load model + loading_info = None + + # Keep in fp32 modules + keep_in_fp32_modules = None + use_keep_in_fp32_modules = False + + if pretrained_model_name_or_path is not None: + pretrained_model_name_or_path = str(pretrained_model_name_or_path) + is_local = os.path.isdir(pretrained_model_name_or_path) + if is_local: + if from_tf and os.path.isfile( + os.path.join(pretrained_model_name_or_path, subfolder, TF_WEIGHTS_NAME + ".index") + ): + # Load from a TF 1.0 checkpoint in priority if from_tf + archive_file = os.path.join(pretrained_model_name_or_path, subfolder, TF_WEIGHTS_NAME + ".index") + elif from_tf and os.path.isfile( + os.path.join(pretrained_model_name_or_path, subfolder, TF2_WEIGHTS_NAME) + ): + # Load from a TF 2.0 checkpoint in priority if from_tf + archive_file = os.path.join(pretrained_model_name_or_path, subfolder, TF2_WEIGHTS_NAME) + elif from_flax and os.path.isfile( + os.path.join(pretrained_model_name_or_path, subfolder, FLAX_WEIGHTS_NAME) + ): + # Load from a Flax checkpoint in priority if from_flax + archive_file = os.path.join(pretrained_model_name_or_path, subfolder, FLAX_WEIGHTS_NAME) + elif use_safetensors is not False and os.path.isfile( + os.path.join(pretrained_model_name_or_path, subfolder, _add_variant(SAFE_WEIGHTS_NAME, variant)) + ): + # Load from a safetensors checkpoint + archive_file = os.path.join( + pretrained_model_name_or_path, subfolder, _add_variant(SAFE_WEIGHTS_NAME, variant) + ) + elif use_safetensors is not False and os.path.isfile( + os.path.join( + pretrained_model_name_or_path, subfolder, _add_variant(SAFE_WEIGHTS_INDEX_NAME, variant) + ) + ): + # Load from a sharded safetensors checkpoint + archive_file = os.path.join( + pretrained_model_name_or_path, subfolder, _add_variant(SAFE_WEIGHTS_INDEX_NAME, variant) + ) + is_sharded = True + elif os.path.isfile( + os.path.join(pretrained_model_name_or_path, subfolder, _add_variant(WEIGHTS_NAME, variant)) + ): + # Load from a PyTorch checkpoint + archive_file = os.path.join( + pretrained_model_name_or_path, subfolder, _add_variant(WEIGHTS_NAME, variant) + ) + elif os.path.isfile( + os.path.join(pretrained_model_name_or_path, subfolder, + _add_variant(WEIGHTS_INDEX_NAME, variant)) + ): + # Load from a sharded PyTorch checkpoint + archive_file = os.path.join( + pretrained_model_name_or_path, subfolder, _add_variant(WEIGHTS_INDEX_NAME, variant) + ) + is_sharded = True + # At this stage we don't have a weight file so we will raise an error. + elif os.path.isfile( + os.path.join(pretrained_model_name_or_path, subfolder, TF_WEIGHTS_NAME + ".index") + ) or os.path.isfile(os.path.join(pretrained_model_name_or_path, subfolder, TF2_WEIGHTS_NAME)): + raise EnvironmentError( + f"Error no file named {_add_variant(WEIGHTS_NAME, variant)} found in directory" + f" {pretrained_model_name_or_path} but there is a file for TensorFlow weights. Use" + " `from_tf=True` to load this model from those weights." + ) + elif os.path.isfile(os.path.join(pretrained_model_name_or_path, subfolder, FLAX_WEIGHTS_NAME)): + raise EnvironmentError( + f"Error no file named {_add_variant(WEIGHTS_NAME, variant)} found in directory" + f" {pretrained_model_name_or_path} but there is a file for Flax weights. Use `from_flax=True`" + " to load this model from those weights." + ) + else: + raise EnvironmentError( + f"Error no file named {_add_variant(WEIGHTS_NAME, variant)}, {TF2_WEIGHTS_NAME}," + f" {TF_WEIGHTS_NAME + '.index'} or {FLAX_WEIGHTS_NAME} found in directory" + f" {pretrained_model_name_or_path}." + ) + elif os.path.isfile(os.path.join(subfolder, pretrained_model_name_or_path)): + archive_file = pretrained_model_name_or_path + is_local = True + elif os.path.isfile(os.path.join(subfolder, pretrained_model_name_or_path + ".index")): + if not from_tf: + raise ValueError( + f"We found a TensorFlow checkpoint at {pretrained_model_name_or_path + '.index'}, please set " + "from_tf to True to load from this checkpoint." + ) + archive_file = os.path.join(subfolder, pretrained_model_name_or_path + ".index") + is_local = True + elif is_remote_url(pretrained_model_name_or_path): + filename = pretrained_model_name_or_path + resolved_archive_file = download_url(pretrained_model_name_or_path) + else: + # set correct filename + if from_tf: + filename = TF2_WEIGHTS_NAME + elif from_flax: + filename = FLAX_WEIGHTS_NAME + elif use_safetensors is not False: + filename = _add_variant(SAFE_WEIGHTS_NAME, variant) + else: + filename = _add_variant(WEIGHTS_NAME, variant) + + try: + # Load from URL or cache if already cached + cached_file_kwargs = { + "cache_dir": cache_dir, + "force_download": force_download, + "proxies": proxies, + "resume_download": resume_download, + "local_files_only": local_files_only, + "use_auth_token": use_auth_token, + "user_agent": user_agent, + "revision": revision, + "subfolder": subfolder, + "_raise_exceptions_for_missing_entries": False, + "_commit_hash": commit_hash, + } + resolved_archive_file = cached_file(pretrained_model_name_or_path, filename, **cached_file_kwargs) + + # Since we set _raise_exceptions_for_missing_entries=False, we don't get an exception but a None + # result when internet is up, the repo and revision exist, but the file does not. + if resolved_archive_file is None and filename == _add_variant(SAFE_WEIGHTS_NAME, variant): + # Maybe the checkpoint is sharded, we try to grab the index name in this case. + resolved_archive_file = cached_file( + pretrained_model_name_or_path, + _add_variant(SAFE_WEIGHTS_INDEX_NAME, variant), + **cached_file_kwargs, + ) + if resolved_archive_file is not None: + is_sharded = True + elif use_safetensors: + raise EnvironmentError( + f" {_add_variant(SAFE_WEIGHTS_NAME, variant)} or {_add_variant(SAFE_WEIGHTS_INDEX_NAME, variant)} and thus cannot be loaded with `safetensors`. Please make sure that the model has been saved with `safe_serialization=True` or do not set `use_safetensors=True`." + ) + else: + # This repo has no safetensors file of any kind, we switch to PyTorch. + filename = _add_variant(WEIGHTS_NAME, variant) + resolved_archive_file = cached_file( + pretrained_model_name_or_path, filename, **cached_file_kwargs + ) + if resolved_archive_file is None and filename == _add_variant(WEIGHTS_NAME, variant): + # Maybe the checkpoint is sharded, we try to grab the index name in this case. + resolved_archive_file = cached_file( + pretrained_model_name_or_path, + _add_variant(WEIGHTS_INDEX_NAME, variant), + **cached_file_kwargs, + ) + if resolved_archive_file is not None: + is_sharded = True + if resolved_archive_file is None: + # Otherwise, maybe there is a TF or Flax model file. We try those to give a helpful error + # message. + has_file_kwargs = { + "revision": revision, + "proxies": proxies, + "use_auth_token": use_auth_token, + } + if has_file(pretrained_model_name_or_path, TF2_WEIGHTS_NAME, **has_file_kwargs): + raise EnvironmentError( + f"{pretrained_model_name_or_path} does not appear to have a file named" + f" {_add_variant(WEIGHTS_NAME, variant)} but there is a file for TensorFlow weights." + " Use `from_tf=True` to load this model from those weights." + ) + elif has_file(pretrained_model_name_or_path, FLAX_WEIGHTS_NAME, **has_file_kwargs): + raise EnvironmentError( + f"{pretrained_model_name_or_path} does not appear to have a file named" + f" {_add_variant(WEIGHTS_NAME, variant)} but there is a file for Flax weights. Use" + " `from_flax=True` to load this model from those weights." + ) + elif variant is not None and has_file( + pretrained_model_name_or_path, WEIGHTS_NAME, **has_file_kwargs + ): + raise EnvironmentError( + f"{pretrained_model_name_or_path} does not appear to have a file named" + f" {_add_variant(WEIGHTS_NAME, variant)} but there is a file without the variant" + f" {variant}. Use `variant=None` to load this model from those weights." + ) + else: + raise EnvironmentError( + f"{pretrained_model_name_or_path} does not appear to have a file named" + f" {_add_variant(WEIGHTS_NAME, variant)}, {TF2_WEIGHTS_NAME}, {TF_WEIGHTS_NAME} or" + f" {FLAX_WEIGHTS_NAME}." + ) + except EnvironmentError: + # Raise any environment error raise by `cached_file`. It will have a helpful error message adapted + # to the original exception. + raise + except Exception: + # For any other exception, we throw a generic error. + raise EnvironmentError( + f"Can't load the model for '{pretrained_model_name_or_path}'. If you were trying to load it" + " from 'https://huggingface.co/models', make sure you don't have a local directory with the" + f" same name. Otherwise, make sure '{pretrained_model_name_or_path}' is the correct path to a" + f" directory containing a file named {_add_variant(WEIGHTS_NAME, variant)}," + f" {TF2_WEIGHTS_NAME}, {TF_WEIGHTS_NAME} or {FLAX_WEIGHTS_NAME}." + ) + + if is_local: + logger.info(f"loading weights file {archive_file}") + resolved_archive_file = archive_file + else: + logger.info(f"loading weights file {filename} from cache at {resolved_archive_file}") + else: + resolved_archive_file = None + + # We'll need to download and cache each checkpoint shard if the checkpoint is sharded. + if is_sharded: + # rsolved_archive_file becomes a list of files that point to the different checkpoint shards in this case. + resolved_archive_file, sharded_metadata = get_checkpoint_shard_files( + pretrained_model_name_or_path, + resolved_archive_file, + cache_dir=cache_dir, + force_download=force_download, + proxies=proxies, + resume_download=resume_download, + local_files_only=local_files_only, + use_auth_token=use_auth_token, + user_agent=user_agent, + revision=revision, + subfolder=subfolder, + _commit_hash=commit_hash, + ) + + # load pt weights early so that we know which dtype to init the model under + if from_pt: + if not is_sharded and state_dict is None: + # Time to load the checkpoint + state_dict = load_state_dict(resolved_archive_file) + + # set dtype to instantiate the model under: + # 1. If torch_dtype is not None, we use that dtype + # 2. If torch_dtype is "auto", we auto-detect dtype from the loaded state_dict, by checking its first + # weights entry that is of a floating type - we assume all floating dtype weights are of the same dtype + # we also may have config.torch_dtype available, but we won't rely on it till v5 + dtype_orig = None + + if torch_dtype is not None: + if isinstance(torch_dtype, str): + if torch_dtype == "auto": + if hasattr(config, "torch_dtype") and config.torch_dtype is not None: + torch_dtype = config.torch_dtype + logger.info(f"Will use torch_dtype={torch_dtype} as defined in model's config object") + else: + if is_sharded and "dtype" in sharded_metadata: + torch_dtype = sharded_metadata["dtype"] + elif not is_sharded: + torch_dtype = get_state_dict_dtype(state_dict) + else: + one_state_dict = load_state_dict(resolved_archive_file[0]) + torch_dtype = get_state_dict_dtype(one_state_dict) + del one_state_dict # free CPU memory + logger.info( + "Since the `torch_dtype` attribute can't be found in model's config object, " + "will use torch_dtype={torch_dtype} as derived from model's weights" + ) + else: + raise ValueError( + f'`torch_dtype` can be either `torch.dtype` or `"auto"`, but received {torch_dtype}' + ) + dtype_orig = cls._set_default_torch_dtype(torch_dtype) + + # Check if `_keep_in_fp32_modules` is not None + use_keep_in_fp32_modules = ( + (cls._keep_in_fp32_modules is not None) + and is_accelerate_available() + and (torch_dtype == torch.float16 or load_in_4bit or load_in_8bit) + ) + if ( + (cls._keep_in_fp32_modules is not None) + and not is_accelerate_available() + and torch_dtype == torch.float16 + ): + logger.warning( + "For stability purposes, it is recommended to have accelerate installed when using this model in" + " torch.float16, please install it with `pip install accelerate`" + ) + + if is_sharded: + loaded_state_dict_keys = sharded_metadata["all_checkpoint_keys"] + else: + loaded_state_dict_keys = list(state_dict.keys()) + if low_cpu_mem_usage or use_keep_in_fp32_modules: + state_dict = None + + config.name_or_path = pretrained_model_name_or_path + + # Instantiate model. + init_contexts = [no_init_weights(_enable=_fast_init)] + + if is_deepspeed_zero3_enabled(): + import deepspeed + + logger.info("Detected DeepSpeed ZeRO-3: activating zero.init() for this model") + init_contexts = [deepspeed.zero.Init(config_dict_or_path=deepspeed_config())] + init_contexts + elif load_in_8bit or load_in_4bit or low_cpu_mem_usage: + init_contexts.append(init_empty_weights()) + + with ContextManagers(init_contexts): + model = cls(config, *model_args, **model_kwargs) + + # Check first if we are `from_pt` + if use_keep_in_fp32_modules: + low_cpu_mem_usage = True + keep_in_fp32_modules = model._keep_in_fp32_modules + else: + keep_in_fp32_modules = [] + + if load_in_8bit or load_in_4bit: + from .utils.bitsandbytes import get_keys_to_not_convert, replace_with_bnb_linear + + llm_int8_skip_modules = quantization_config.llm_int8_skip_modules + load_in_8bit_fp32_cpu_offload = quantization_config.llm_int8_enable_fp32_cpu_offload + + logger.info("Detected 8-bit loading: activating 8-bit loading for this model") + + # We keep some modules such as the lm_head in their original dtype for numerical stability reasons + if llm_int8_skip_modules is None: + modules_to_not_convert = get_keys_to_not_convert(model) + else: + modules_to_not_convert = llm_int8_skip_modules + + if not isinstance(modules_to_not_convert, list): + modules_to_not_convert = [modules_to_not_convert] + + modules_to_not_convert.extend(keep_in_fp32_modules) + + # Extend the modules to not convert to keys that are supposed to be offloaded to `cpu` or `disk` + if isinstance(device_map, dict) and len(device_map.keys()) > 1: + keys_on_cpu = [key for key, value in device_map.items() if value in ["disk", "cpu"]] + + if len(keys_on_cpu) > 0 and not load_in_8bit_fp32_cpu_offload: + raise ValueError( + "If you want to offload some keys to `cpu` or `disk`, you need to set " + "`llm_int8_enable_fp32_cpu_offload=True`. Note that these modules will not be " + " converted to 8-bit but kept in 32-bit." + ) + + modules_to_not_convert.extend(keys_on_cpu) + + supports_4bit = version.parse(importlib_metadata.version("bitsandbytes")) >= version.parse("0.39.0") + + if load_in_4bit and not supports_4bit: + raise ValueError( + "You have a version of `bitsandbytes` that is not compatible with 4bit inference and training" + " make sure you have the latest version of `bitsandbytes` installed" + ) + + model = replace_with_bnb_linear( + model, modules_to_not_convert=modules_to_not_convert, quantization_config=quantization_config + ) + # training in 8-bit is only available in 0.37.0+ + model._is_quantized_training_enabled = version.parse( + importlib_metadata.version("bitsandbytes") + ) >= version.parse("0.37.0") + + model.config.quantization_config = quantization_config + model.is_8bit_serializable = is_8bit_serializable + + if load_in_8bit and torch_dtype is None: + logger.warning( + "You are loading your model in 8bit but you did not specify a `torch_dtype` attribute." + "All non-linear modules will be loaded in full precision." + " If you want to load the other modules in other precision, please specify a `torch_dtype` attribute." + ) + + if isinstance(device_map, str): + special_dtypes = {} + if load_in_8bit or load_in_4bit: + special_dtypes.update( + { + name: torch_dtype + for name, _ in model.named_parameters() + if any(m in name for m in modules_to_not_convert) + } + ) + + special_dtypes.update( + { + name: torch.float32 + for name, _ in model.named_parameters() + if any(m in name for m in keep_in_fp32_modules) + } + ) + + target_dtype = torch_dtype + + if load_in_4bit: + if version.parse(importlib_metadata.version("accelerate")) > version.parse("0.19.0"): + from accelerate.utils import CustomDtype + + target_dtype = CustomDtype.INT4 + else: + raise ValueError( + "You are using `device_map='auto'` on a 4bit loaded version of the model. To automatically compute" + " the appropriate device map, you should upgrade your `accelerate` library," + "`pip install --upgrade accelerate` or install it from source to support fp4 auto device map" + "calculation. You may encounter unexpected behavior, or pass your own device map" + ) + elif load_in_8bit: + target_dtype = torch.int8 + + if model._no_split_modules is None: + raise ValueError( + f"{model.__class__.__name__} does not support `device_map='{device_map}'`. To implement support, the model" + "class needs to implement the `_no_split_modules` attribute." + ) + no_split_modules = model._no_split_modules + if device_map not in ["auto", "balanced", "balanced_low_0", "sequential"]: + raise ValueError( + "If passing a string for `device_map`, please choose 'auto', 'balanced', 'balanced_low_0' or " + "'sequential'." + ) + elif device_map in ["balanced", "balanced_low_0"] and get_balanced_memory is None: + raise ValueError(f"`device_map={device_map}` requires a source install of Accelerate.") + + kwargs = {"no_split_module_classes": no_split_modules} + if "special_dtypes" in inspect.signature(infer_auto_device_map).parameters: + kwargs["special_dtypes"] = special_dtypes + elif len(special_dtypes) > 0: + logger.warn( + "This model has some weights that should be kept in higher precision, you need to upgrade " + "`accelerate` to properly deal with them (`pip install --upgrade accelerate`)." + ) + if device_map != "sequential" and get_balanced_memory is not None: + max_memory = get_balanced_memory( + model, + dtype=target_dtype, + low_zero=(device_map == "balanced_low_0"), + max_memory=max_memory, + **kwargs, + ) + kwargs["max_memory"] = max_memory + # Make sure tied weights are tied before creating the device map. + model.tie_weights() + device_map = infer_auto_device_map(model, dtype=target_dtype, **kwargs) + + if load_in_8bit or load_in_4bit: + # The LM head / tied weights or any last module can stay on disk / CPU + device_map_without_lm_head = { + key: device_map[key] for key in device_map.keys() if key not in modules_to_not_convert + } + if "cpu" in device_map_without_lm_head.values() or "disk" in device_map_without_lm_head.values(): + raise ValueError( + """ + Some modules are dispatched on the CPU or the disk. Make sure you have enough GPU RAM to fit + the quantized model. If you want to dispatch the model on the CPU or the disk while keeping + these modules in 32-bit, you need to set `load_in_8bit_fp32_cpu_offload=True` and pass a custom + `device_map` to `from_pretrained`. Check + https://huggingface.co/docs/transformers/main/en/main_classes/quantization#offload-between-cpu-and-gpu + for more details. + """ + ) + del device_map_without_lm_head + + elif device_map is not None: + model.tie_weights() + tied_params = find_tied_parameters(model) + # check if we don't have tied param in different devices + if check_tied_parameters_on_same_device is not None: + check_tied_parameters_on_same_device(tied_params, device_map) + + if from_tf: + if resolved_archive_file.endswith(".index"): + # Load from a TensorFlow 1.X checkpoint - provided by original authors + model = cls.load_tf_weights(model, config, resolved_archive_file[:-6]) # Remove the '.index' + else: + # Load from our TensorFlow 2.0 checkpoints + try: + from .modeling_tf_pytorch_utils import load_tf2_checkpoint_in_pytorch_model + + model, loading_info = load_tf2_checkpoint_in_pytorch_model( + model, resolved_archive_file, allow_missing_keys=True, output_loading_info=True + ) + except ImportError: + logger.error( + "Loading a TensorFlow model in PyTorch, requires both PyTorch and TensorFlow to be installed." + " Please see https://pytorch.org/ and https://www.tensorflow.org/install/ for installation" + " instructions." + ) + raise + elif from_flax: + try: + from .modeling_flax_pytorch_utils import load_flax_checkpoint_in_pytorch_model + + model = load_flax_checkpoint_in_pytorch_model(model, resolved_archive_file) + except ImportError: + logger.error( + "Loading a Flax model in PyTorch, requires both PyTorch and Flax to be installed. Please see" + " https://pytorch.org/ and https://flax.readthedocs.io/en/latest/installation.html for" + " installation instructions." + ) + raise + elif from_pt: + # restore default dtype + if dtype_orig is not None: + torch.set_default_dtype(dtype_orig) + + ( + model, + missing_keys, + unexpected_keys, + mismatched_keys, + offload_index, + error_msgs, + ) = cls._load_pretrained_model( + model, + state_dict, + loaded_state_dict_keys, # XXX: rename? + resolved_archive_file, + pretrained_model_name_or_path, + ignore_mismatched_sizes=ignore_mismatched_sizes, + sharded_metadata=sharded_metadata, + _fast_init=_fast_init, + low_cpu_mem_usage=low_cpu_mem_usage, + device_map=device_map, + offload_folder=offload_folder, + offload_state_dict=offload_state_dict, + dtype=torch_dtype, + is_quantized=(load_in_8bit or load_in_4bit), + keep_in_fp32_modules=keep_in_fp32_modules, + ) + + model.is_loaded_in_4bit = load_in_4bit + model.is_loaded_in_8bit = load_in_8bit + model.is_quantized = load_in_8bit or load_in_4bit + + # make sure token embedding weights are still tied if needed + model.tie_weights() + + # Set model in evaluation mode to deactivate DropOut modules by default + model.eval() + + # If it is a model with generation capabilities, attempt to load the generation config + if model.can_generate(): + try: + model.generation_config = GenerationConfig.from_pretrained( + pretrained_model_name_or_path, + cache_dir=cache_dir, + force_download=force_download, + resume_download=resume_download, + proxies=proxies, + local_files_only=local_files_only, + use_auth_token=use_auth_token, + revision=revision, + subfolder=subfolder, + _from_auto=from_auto_class, + _from_pipeline=from_pipeline, + **kwargs, + ) + except (OSError, TypeError): + logger.info( + "Generation config file not found, using a generation config created from the model config." + ) + pass + + # Dispatch model with hooks on all devices if necessary + if device_map is not None: + kwargs = {"device_map": device_map, "offload_dir": offload_folder, "offload_index": offload_index} + if "skip_keys" in inspect.signature(dispatch_model).parameters: + kwargs["skip_keys"] = model._skip_keys_device_placement + dispatch_model(model, **kwargs) + + if output_loading_info: + if loading_info is None: + loading_info = { + "missing_keys": missing_keys, + "unexpected_keys": unexpected_keys, + "mismatched_keys": mismatched_keys, + "error_msgs": error_msgs, + } + return model, loading_info + + return model + + @classmethod + def _load_pretrained_model( + cls, + model, + state_dict, + loaded_keys, + resolved_archive_file, + pretrained_model_name_or_path, + ignore_mismatched_sizes=False, + sharded_metadata=None, + _fast_init=True, + low_cpu_mem_usage=False, + device_map=None, + offload_folder=None, + offload_state_dict=None, + dtype=None, + is_quantized=False, + keep_in_fp32_modules=None, + ): + is_safetensors = False + if is_quantized: + from .utils.bitsandbytes import set_module_quantized_tensor_to_device + + if device_map is not None and "disk" in device_map.values(): + archive_file = ( + resolved_archive_file[0] if isinstance(resolved_archive_file, (list, tuple)) else resolved_archive_file + ) + is_safetensors = archive_file.endswith(".safetensors") + if offload_folder is None and not is_safetensors: + raise ValueError( + "The current `device_map` had weights offloaded to the disk. Please provide an `offload_folder`" + " for them. Alternatively, make sure you have `safetensors` installed if the model you are using" + " offers the weights in this format." + ) + if offload_folder is not None: + os.makedirs(offload_folder, exist_ok=True) + if offload_state_dict is None: + offload_state_dict = True + + is_sharded_safetensors = is_safetensors and sharded_metadata is not None + # Retrieve missing & unexpected_keys + model_state_dict = model.state_dict() + expected_keys = list(model_state_dict.keys()) + prefix = model.base_model_prefix + + def _fix_key(key): + if "beta" in key: + return key.replace("beta", "bias") + if "gamma" in key: + return key.replace("gamma", "weight") + return key + + original_loaded_keys = loaded_keys + loaded_keys = [_fix_key(key) for key in loaded_keys] + + if len(prefix) > 0: + has_prefix_module = any(s.startswith(prefix) for s in loaded_keys) + expects_prefix_module = any(s.startswith(prefix) for s in expected_keys) + else: + has_prefix_module = False + expects_prefix_module = False + + # key re-naming operations are never done on the keys + # that are loaded, but always on the keys of the newly initialized model + remove_prefix_from_model = not has_prefix_module and expects_prefix_module + add_prefix_to_model = has_prefix_module and not expects_prefix_module + + if remove_prefix_from_model: + _prefix = f"{prefix}." + expected_keys_not_prefixed = [s for s in expected_keys if not s.startswith(_prefix)] + expected_keys = [s[len(_prefix):] if s.startswith(_prefix) else s for s in expected_keys] + elif add_prefix_to_model: + expected_keys = [".".join([prefix, s]) for s in expected_keys] + + missing_keys = list(set(expected_keys) - set(loaded_keys)) + unexpected_keys = list(set(loaded_keys) - set(expected_keys)) + + if find_tied_parameters is not None: + model.tie_weights() + tied_params = find_tied_parameters(model) + else: + tied_params = [] + _missing = [] + for k in missing_keys: + found = False + for group in tied_params: + if k in group: + found = True + if len(group) > 2: + group.remove(k) + else: + _missing.append(k) + if not found: + _missing.append(k) + missing_keys = _missing + + # Some models may have keys that are not in the state by design, removing them before needlessly warning + # the user. + if cls._keys_to_ignore_on_load_missing is not None: + for pat in cls._keys_to_ignore_on_load_missing: + missing_keys = [k for k in missing_keys if re.search(pat, k) is None] + + if cls._keys_to_ignore_on_load_unexpected is not None: + for pat in cls._keys_to_ignore_on_load_unexpected: + unexpected_keys = [k for k in unexpected_keys if re.search(pat, k) is None] + + # retrieve weights on meta device and put them back on CPU. + # This is not ideal in terms of memory, but if we don't do that not, we can't initialize them in the next step + if low_cpu_mem_usage: + for key in missing_keys: + if key in list(model_state_dict.keys()): + key = key + elif f"{prefix}.{key}" in list(model_state_dict.keys()): + key = f"{prefix}.{key}" + elif key.startswith(prefix) and ".".join(key.split(".")[1:]) in list(model_state_dict.keys()): + key = ".".join(key.split(".")[1:]) + param = model_state_dict[key] + + # upcast in fp32 if any + target_dtype = dtype + if ( + keep_in_fp32_modules is not None + and dtype == torch.float16 + and any(module_to_keep_in_fp32 in key for module_to_keep_in_fp32 in keep_in_fp32_modules) + ): + target_dtype = torch.float32 + + if param.device == torch.device("meta"): + if not (is_quantized): + set_module_tensor_to_device(model, key, "cpu", torch.empty(*param.size(), dtype=target_dtype)) + else: + set_module_quantized_tensor_to_device( + model, key, "cpu", torch.empty(*param.size(), dtype=target_dtype) + ) + + # retrieve unintialized modules and initialize before maybe overriding that with the pretrained weights. + if _fast_init: + if remove_prefix_from_model: + _loaded_keys = [f"{prefix}.{k}" for k in loaded_keys] + elif add_prefix_to_model: + _loaded_keys = [k[len(prefix) + 1:] for k in loaded_keys] + else: + _loaded_keys = loaded_keys + set_initialized_submodules(model, _loaded_keys) + # This will only initialize submodules that are not marked as initialized by the line above. + model.apply(model._initialize_weights) + + # Set some modules to fp32 if any + if keep_in_fp32_modules is not None: + for name, param in model.named_parameters(): + if any(module_to_keep_in_fp32 in name for module_to_keep_in_fp32 in keep_in_fp32_modules): + param = param.to(torch.float32) + + # Make sure we are able to load base models as well as derived models (with heads) + start_prefix = "" + model_to_load = model + if len(cls.base_model_prefix) > 0 and not hasattr(model, cls.base_model_prefix) and has_prefix_module: + start_prefix = cls.base_model_prefix + "." + if len(cls.base_model_prefix) > 0 and hasattr(model, cls.base_model_prefix) and not has_prefix_module: + model_to_load = getattr(model, cls.base_model_prefix) + base_model_expected_keys = list(model_to_load.state_dict().keys()) + if any(key in expected_keys_not_prefixed and key not in base_model_expected_keys for key in loaded_keys): + raise ValueError( + "The state dictionary of the model you are trying to load is corrupted. Are you sure it was " + "properly saved?" + ) + if device_map is not None: + device_map = {k.replace(f"{cls.base_model_prefix}.", ""): v for k, v in device_map.items()} + + def _find_mismatched_keys( + state_dict, + model_state_dict, + loaded_keys, + add_prefix_to_model, + remove_prefix_from_model, + ignore_mismatched_sizes, + ): + mismatched_keys = [] + if ignore_mismatched_sizes: + for checkpoint_key in loaded_keys: + model_key = checkpoint_key + if remove_prefix_from_model: + # The model key starts with `prefix` but `checkpoint_key` doesn't so we add it. + model_key = f"{prefix}.{checkpoint_key}" + elif add_prefix_to_model: + # The model key doesn't start with `prefix` but `checkpoint_key` does so we remove it. + model_key = ".".join(checkpoint_key.split(".")[1:]) + + if ( + model_key in model_state_dict + and state_dict[checkpoint_key].shape != model_state_dict[model_key].shape + ): + mismatched_keys.append( + (checkpoint_key, state_dict[checkpoint_key].shape, model_state_dict[model_key].shape) + ) + del state_dict[checkpoint_key] + return mismatched_keys + + if resolved_archive_file is not None: + folder = os.path.sep.join(resolved_archive_file[0].split(os.path.sep)[:-1]) + else: + folder = None + if device_map is not None and is_safetensors: + param_device_map = expand_device_map(device_map, original_loaded_keys) + + str_dtype = str(dtype).replace("torch.", "") if dtype is not None else "float32" + if sharded_metadata is None: + archive_file = ( + resolved_archive_file[0] + if isinstance(resolved_archive_file, (list, tuple)) + else resolved_archive_file + ) + weight_map = {p: archive_file for p in original_loaded_keys} + else: + weight_map = {p: os.path.join(folder, f) for p, f in sharded_metadata["weight_map"].items()} + offload_index = { + p: {"safetensors_file": f, "weight_name": p, "dtype": str_dtype} + for p, f in weight_map.items() + if param_device_map[p] == "disk" + } + + if state_dict is not None: + # Whole checkpoint + mismatched_keys = _find_mismatched_keys( + state_dict, + model_state_dict, + original_loaded_keys, + add_prefix_to_model, + remove_prefix_from_model, + ignore_mismatched_sizes, + ) + error_msgs = _load_state_dict_into_model(model_to_load, state_dict, start_prefix) + offload_index = None + else: + # Sharded checkpoint or whole but low_cpu_mem_usage==True + + # This should always be a list but, just to be sure. + if not isinstance(resolved_archive_file, list): + resolved_archive_file = [resolved_archive_file] + + error_msgs = [] + mismatched_keys = [] + if not is_safetensors: + offload_index = {} if device_map is not None and "disk" in device_map.values() else None + if offload_state_dict: + state_dict_folder = tempfile.mkdtemp() + state_dict_index = {} + else: + state_dict_folder = None + state_dict_index = None + + if is_sharded_safetensors: + disk_only_shard_files = get_disk_only_shard_files(device_map, sharded_metadata=sharded_metadata) + disk_only_shard_files = [os.path.join(folder, f) for f in disk_only_shard_files] + else: + disk_only_shard_files = [] + + if len(resolved_archive_file) > 1: + resolved_archive_file = logging.tqdm(resolved_archive_file, desc="Loading checkpoint shards") + for shard_file in resolved_archive_file: + # Skip the load for shards that only contain disk-offloaded weights when using safetensors for the offload. + if shard_file in disk_only_shard_files: + continue + state_dict = load_state_dict(shard_file) + + # Mistmatched keys contains tuples key/shape1/shape2 of weights in the checkpoint that have a shape not + # matching the weights in the model. + mismatched_keys += _find_mismatched_keys( + state_dict, + model_state_dict, + original_loaded_keys, + add_prefix_to_model, + remove_prefix_from_model, + ignore_mismatched_sizes, + ) + + if low_cpu_mem_usage: + new_error_msgs, offload_index, state_dict_index = _load_state_dict_into_meta_model( + model_to_load, + state_dict, + loaded_keys, + start_prefix, + expected_keys, + device_map=device_map, + offload_folder=offload_folder, + offload_index=offload_index, + state_dict_folder=state_dict_folder, + state_dict_index=state_dict_index, + dtype=dtype, + is_quantized=is_quantized, + is_safetensors=is_safetensors, + keep_in_fp32_modules=keep_in_fp32_modules, + ) + error_msgs += new_error_msgs + else: + error_msgs += _load_state_dict_into_model(model_to_load, state_dict, start_prefix) + + # force memory release + del state_dict + gc.collect() + + if offload_index is not None and len(offload_index) > 0: + if model != model_to_load: + # We need to add the prefix of the base model + prefix = cls.base_model_prefix + if not is_safetensors: + for weight_name in offload_index: + shutil.move( + os.path.join(offload_folder, f"{weight_name}.dat"), + os.path.join(offload_folder, f"{prefix}.{weight_name}.dat"), + ) + offload_index = {f"{prefix}.{key}": value for key, value in offload_index.items()} + if not is_safetensors: + save_offload_index(offload_index, offload_folder) + offload_index = None + + if offload_state_dict: + # Load back temporarily offloaded state dict + load_offloaded_weights(model_to_load, state_dict_index, state_dict_folder) + shutil.rmtree(state_dict_folder) + + if len(error_msgs) > 0: + error_msg = "\n\t".join(error_msgs) + if "size mismatch" in error_msg: + error_msg += ( + "\n\tYou may consider adding `ignore_mismatched_sizes=True` in the model `from_pretrained` method." + ) + raise RuntimeError(f"Error(s) in loading state_dict for {model.__class__.__name__}:\n\t{error_msg}") + + if is_quantized: + unexpected_keys = [elem for elem in unexpected_keys if "SCB" not in elem] + missing_keys = [elem for elem in missing_keys if "SCB" not in elem] + + if len(unexpected_keys) > 0: + logger.warning( + f"Some weights of the model checkpoint at {pretrained_model_name_or_path} were not used when" + f" initializing {model.__class__.__name__}: {unexpected_keys}\n- This IS expected if you are" + f" initializing {model.__class__.__name__} from the checkpoint of a model trained on another task or" + " with another architecture (e.g. initializing a BertForSequenceClassification model from a" + " BertForPreTraining model).\n- This IS NOT expected if you are initializing" + f" {model.__class__.__name__} from the checkpoint of a model that you expect to be exactly identical" + " (initializing a BertForSequenceClassification model from a BertForSequenceClassification model)." + ) + else: + logger.info(f"All model checkpoint weights were used when initializing {model.__class__.__name__}.\n") + if len(missing_keys) > 0: + logger.warning( + f"Some weights of {model.__class__.__name__} were not initialized from the model checkpoint at" + f" {pretrained_model_name_or_path} and are newly initialized: {missing_keys}\nYou should probably" + " TRAIN this model on a down-stream task to be able to use it for predictions and inference." + ) + elif len(mismatched_keys) == 0: + logger.info( + f"All the weights of {model.__class__.__name__} were initialized from the model checkpoint at" + f" {pretrained_model_name_or_path}.\nIf your task is similar to the task the model of the checkpoint" + f" was trained on, you can already use {model.__class__.__name__} for predictions without further" + " training." + ) + if len(mismatched_keys) > 0: + mismatched_warning = "\n".join( + [ + f"- {key}: found shape {shape1} in the checkpoint and {shape2} in the model instantiated" + for key, shape1, shape2 in mismatched_keys + ] + ) + logger.warning( + f"Some weights of {model.__class__.__name__} were not initialized from the model checkpoint at" + f" {pretrained_model_name_or_path} and are newly initialized because the shapes did not" + f" match:\n{mismatched_warning}\nYou should probably TRAIN this model on a down-stream task to be able" + " to use it for predictions and inference." + ) + + return model, missing_keys, unexpected_keys, mismatched_keys, offload_index, error_msgs + + def retrieve_modules_from_names(self, names, add_prefix=False, remove_prefix=False): + module_keys = {".".join(key.split(".")[:-1]) for key in names} + + # torch.nn.ParameterList is a special case where two parameter keywords + # are appended to the module name, *e.g.* bert.special_embeddings.0 + module_keys = module_keys.union( + {".".join(key.split(".")[:-2]) for key in names if len(key) > 0 and key[-1].isdigit()} + ) + + retrieved_modules = [] + # retrieve all modules that has at least one missing weight name + for name, module in self.named_modules(): + if remove_prefix: + _prefix = f"{self.base_model_prefix}." + name = name[len(_prefix):] if name.startswith(_prefix) else name + elif add_prefix: + name = ".".join([self.base_model_prefix, name]) if len(name) > 0 else self.base_model_prefix + + if name in module_keys: + retrieved_modules.append(module) + + return retrieved_modules + + @staticmethod + def _load_pretrained_model_low_mem(model, loaded_state_dict_keys, resolved_archive_file, start_prefix=""): + """ + This is an experimental function that loads the model using ~1.x model size CPU memory + + Before you call it do: + + 1. save which state_dict keys are available + 2. drop state_dict before model is created, since the latter takes 1x model size memory + + Here then we continue: + + 3. switch to the meta device all params/buffers that are going to be replaced from the loaded state_dict + 4. load state_dict 2nd time + 5. replace the params/buffers from the state_dict + + Currently, it doesn't handle missing_keys, unexpected_keys, mismatched_keys. It can't handle deepspeed. + """ + + _move_model_to_meta(model, loaded_state_dict_keys, start_prefix) + state_dict = load_state_dict(resolved_archive_file) + error_msgs = _load_state_dict_into_meta_model(model, state_dict, loaded_state_dict_keys, start_prefix) + return error_msgs + + @classmethod + def register_for_auto_class(cls, auto_class="AutoModel"): + """ + Register this class with a given auto class. This should only be used for custom models as the ones in the + library are already mapped with an auto class. + + + + This API is experimental and may have some slight breaking changes in the next releases. + + + + Args: + auto_class (`str` or `type`, *optional*, defaults to `"AutoModel"`): + The auto class to register this new model with. + """ + if not isinstance(auto_class, str): + auto_class = auto_class.__name__ + + import transformers.models.auto as auto_module + + if not hasattr(auto_module, auto_class): + raise ValueError(f"{auto_class} is not a valid auto class.") + + cls._auto_class = auto_class + + def to_bettertransformer(self) -> "PreTrainedModel": + """ + Converts the model to use [PyTorch's native attention + implementation](https://pytorch.org/docs/stable/generated/torch.nn.MultiheadAttention.html), integrated to + Transformers through [Optimum library](https://huggingface.co/docs/optimum/bettertransformer/overview). Only a + subset of all Transformers models are supported. + + PyTorch's attention fastpath allows to speed up inference through kernel fusions and the use of [nested + tensors](https://pytorch.org/docs/stable/nested.html). Detailed benchmarks can be found in [this blog + post](https://medium.com/pytorch/bettertransformer-out-of-the-box-performance-for-huggingface-transformers-3fbe27d50ab2). + + Returns: + [`PreTrainedModel`]: The model converted to BetterTransformer. + """ + if not is_optimum_available(): + raise ImportError("The package `optimum` is required to use Better Transformer.") + + from optimum.version import __version__ as optimum_version + + if version.parse(optimum_version) < version.parse("1.7.0"): + raise ImportError( + f"Please install optimum>=1.7.0 to use Better Transformer. The version {optimum_version} was found." + ) + + from optimum.bettertransformer import BetterTransformer + + return BetterTransformer.transform(self) + + def reverse_bettertransformer(self): + """ + Reverts the transformation from [`~PreTrainedModel.to_bettertransformer`] so that the original modeling is + used, for example in order to save the model. + + Returns: + [`PreTrainedModel`]: The model converted back to the original modeling. + """ + if not is_optimum_available(): + raise ImportError("The package `optimum` is required to use Better Transformer.") + + from optimum.version import __version__ as optimum_version + + if version.parse(optimum_version) < version.parse("1.7.0"): + raise ImportError( + f"Please install optimum>=1.7.0 to use Better Transformer. The version {optimum_version} was found." + ) + + from optimum.bettertransformer import BetterTransformer + + return BetterTransformer.reverse(self) + + +PreTrainedModel.push_to_hub = copy_func(PreTrainedModel.push_to_hub) +if PreTrainedModel.push_to_hub.__doc__ is not None: + PreTrainedModel.push_to_hub.__doc__ = PreTrainedModel.push_to_hub.__doc__.format( + object="model", object_class="AutoModel", object_files="model file" + ) + + +class PoolerStartLogits(nn.Module): + """ + Compute SQuAD start logits from sequence hidden states. + + Args: + config ([`PretrainedConfig`]): + The config used by the model, will be used to grab the `hidden_size` of the model. + """ + + def __init__(self, config: PretrainedConfig): + super().__init__() + self.dense = nn.Linear(config.hidden_size, 1) + + def forward( + self, hidden_states: torch.FloatTensor, p_mask: Optional[torch.FloatTensor] = None + ) -> torch.FloatTensor: + """ + Args: + hidden_states (`torch.FloatTensor` of shape `(batch_size, seq_len, hidden_size)`): + The final hidden states of the model. + p_mask (`torch.FloatTensor` of shape `(batch_size, seq_len)`, *optional*): + Mask for tokens at invalid position, such as query and special symbols (PAD, SEP, CLS). 1.0 means token + should be masked. + + Returns: + `torch.FloatTensor`: The start logits for SQuAD. + """ + x = self.dense(hidden_states).squeeze(-1) + + if p_mask is not None: + if get_parameter_dtype(self) == torch.float16: + x = x * (1 - p_mask) - 65500 * p_mask + else: + x = x * (1 - p_mask) - 1e30 * p_mask + + return x + + +class PoolerEndLogits(nn.Module): + """ + Compute SQuAD end logits from sequence hidden states. + + Args: + config ([`PretrainedConfig`]): + The config used by the model, will be used to grab the `hidden_size` of the model and the `layer_norm_eps` + to use. + """ + + def __init__(self, config: PretrainedConfig): + super().__init__() + self.dense_0 = nn.Linear(config.hidden_size * 2, config.hidden_size) + self.activation = nn.Tanh() + self.LayerNorm = nn.LayerNorm(config.hidden_size, eps=config.layer_norm_eps) + self.dense_1 = nn.Linear(config.hidden_size, 1) + + def forward( + self, + hidden_states: torch.FloatTensor, + start_states: Optional[torch.FloatTensor] = None, + start_positions: Optional[torch.LongTensor] = None, + p_mask: Optional[torch.FloatTensor] = None, + ) -> torch.FloatTensor: + """ + Args: + hidden_states (`torch.FloatTensor` of shape `(batch_size, seq_len, hidden_size)`): + The final hidden states of the model. + start_states (`torch.FloatTensor` of shape `(batch_size, seq_len, hidden_size)`, *optional*): + The hidden states of the first tokens for the labeled span. + start_positions (`torch.LongTensor` of shape `(batch_size,)`, *optional*): + The position of the first token for the labeled span. + p_mask (`torch.FloatTensor` of shape `(batch_size, seq_len)`, *optional*): + Mask for tokens at invalid position, such as query and special symbols (PAD, SEP, CLS). 1.0 means token + should be masked. + + + + One of `start_states` or `start_positions` should be not `None`. If both are set, `start_positions` overrides + `start_states`. + + + + Returns: + `torch.FloatTensor`: The end logits for SQuAD. + """ + assert ( + start_states is not None or start_positions is not None + ), "One of start_states, start_positions should be not None" + if start_positions is not None: + slen, hsz = hidden_states.shape[-2:] + start_positions = start_positions[:, None, None].expand(-1, -1, hsz) # shape (bsz, 1, hsz) + start_states = hidden_states.gather(-2, start_positions) # shape (bsz, 1, hsz) + start_states = start_states.expand(-1, slen, -1) # shape (bsz, slen, hsz) + + x = self.dense_0(torch.cat([hidden_states, start_states], dim=-1)) + x = self.activation(x) + x = self.LayerNorm(x) + x = self.dense_1(x).squeeze(-1) + + if p_mask is not None: + if get_parameter_dtype(self) == torch.float16: + x = x * (1 - p_mask) - 65500 * p_mask + else: + x = x * (1 - p_mask) - 1e30 * p_mask + + return x + + +class PoolerAnswerClass(nn.Module): + """ + Compute SQuAD 2.0 answer class from classification and start tokens hidden states. + + Args: + config ([`PretrainedConfig`]): + The config used by the model, will be used to grab the `hidden_size` of the model. + """ + + def __init__(self, config): + super().__init__() + self.dense_0 = nn.Linear(config.hidden_size * 2, config.hidden_size) + self.activation = nn.Tanh() + self.dense_1 = nn.Linear(config.hidden_size, 1, bias=False) + + def forward( + self, + hidden_states: torch.FloatTensor, + start_states: Optional[torch.FloatTensor] = None, + start_positions: Optional[torch.LongTensor] = None, + cls_index: Optional[torch.LongTensor] = None, + ) -> torch.FloatTensor: + """ + Args: + hidden_states (`torch.FloatTensor` of shape `(batch_size, seq_len, hidden_size)`): + The final hidden states of the model. + start_states (`torch.FloatTensor` of shape `(batch_size, seq_len, hidden_size)`, *optional*): + The hidden states of the first tokens for the labeled span. + start_positions (`torch.LongTensor` of shape `(batch_size,)`, *optional*): + The position of the first token for the labeled span. + cls_index (`torch.LongTensor` of shape `(batch_size,)`, *optional*): + Position of the CLS token for each sentence in the batch. If `None`, takes the last token. + + + + One of `start_states` or `start_positions` should be not `None`. If both are set, `start_positions` overrides + `start_states`. + + + + Returns: + `torch.FloatTensor`: The SQuAD 2.0 answer class. + """ + # No dependency on end_feature so that we can obtain one single `cls_logits` for each sample. + hsz = hidden_states.shape[-1] + assert ( + start_states is not None or start_positions is not None + ), "One of start_states, start_positions should be not None" + if start_positions is not None: + start_positions = start_positions[:, None, None].expand(-1, -1, hsz) # shape (bsz, 1, hsz) + start_states = hidden_states.gather(-2, start_positions).squeeze(-2) # shape (bsz, hsz) + + if cls_index is not None: + cls_index = cls_index[:, None, None].expand(-1, -1, hsz) # shape (bsz, 1, hsz) + cls_token_state = hidden_states.gather(-2, cls_index).squeeze(-2) # shape (bsz, hsz) + else: + cls_token_state = hidden_states[:, -1, :] # shape (bsz, hsz) + + x = self.dense_0(torch.cat([start_states, cls_token_state], dim=-1)) + x = self.activation(x) + x = self.dense_1(x).squeeze(-1) + + return x + + +@dataclass +class SquadHeadOutput(ModelOutput): + """ + Base class for outputs of question answering models using a [`~modeling_utils.SQuADHead`]. + + Args: + loss (`torch.FloatTensor` of shape `(1,)`, *optional*, returned if both `start_positions` and `end_positions` are provided): + Classification loss as the sum of start token, end token (and is_impossible if provided) classification + losses. + start_top_log_probs (`torch.FloatTensor` of shape `(batch_size, config.start_n_top)`, *optional*, returned if `start_positions` or `end_positions` is not provided): + Log probabilities for the top config.start_n_top start token possibilities (beam-search). + start_top_index (`torch.LongTensor` of shape `(batch_size, config.start_n_top)`, *optional*, returned if `start_positions` or `end_positions` is not provided): + Indices for the top config.start_n_top start token possibilities (beam-search). + end_top_log_probs (`torch.FloatTensor` of shape `(batch_size, config.start_n_top * config.end_n_top)`, *optional*, returned if `start_positions` or `end_positions` is not provided): + Log probabilities for the top `config.start_n_top * config.end_n_top` end token possibilities + (beam-search). + end_top_index (`torch.LongTensor` of shape `(batch_size, config.start_n_top * config.end_n_top)`, *optional*, returned if `start_positions` or `end_positions` is not provided): + Indices for the top `config.start_n_top * config.end_n_top` end token possibilities (beam-search). + cls_logits (`torch.FloatTensor` of shape `(batch_size,)`, *optional*, returned if `start_positions` or `end_positions` is not provided): + Log probabilities for the `is_impossible` label of the answers. + + """ + + loss: Optional[torch.FloatTensor] = None + start_top_log_probs: Optional[torch.FloatTensor] = None + start_top_index: Optional[torch.LongTensor] = None + end_top_log_probs: Optional[torch.FloatTensor] = None + end_top_index: Optional[torch.LongTensor] = None + cls_logits: Optional[torch.FloatTensor] = None + + +class SQuADHead(nn.Module): + r""" + A SQuAD head inspired by XLNet. + + Args: + config ([`PretrainedConfig`]): + The config used by the model, will be used to grab the `hidden_size` of the model and the `layer_norm_eps` + to use. + """ + + def __init__(self, config): + super().__init__() + self.start_n_top = config.start_n_top + self.end_n_top = config.end_n_top + + self.start_logits = PoolerStartLogits(config) + self.end_logits = PoolerEndLogits(config) + self.answer_class = PoolerAnswerClass(config) + + @replace_return_docstrings(output_type=SquadHeadOutput, config_class=PretrainedConfig) + def forward( + self, + hidden_states: torch.FloatTensor, + start_positions: Optional[torch.LongTensor] = None, + end_positions: Optional[torch.LongTensor] = None, + cls_index: Optional[torch.LongTensor] = None, + is_impossible: Optional[torch.LongTensor] = None, + p_mask: Optional[torch.FloatTensor] = None, + return_dict: bool = False, + ) -> Union[SquadHeadOutput, Tuple[torch.FloatTensor]]: + """ + Args: + hidden_states (`torch.FloatTensor` of shape `(batch_size, seq_len, hidden_size)`): + Final hidden states of the model on the sequence tokens. + start_positions (`torch.LongTensor` of shape `(batch_size,)`, *optional*): + Positions of the first token for the labeled span. + end_positions (`torch.LongTensor` of shape `(batch_size,)`, *optional*): + Positions of the last token for the labeled span. + cls_index (`torch.LongTensor` of shape `(batch_size,)`, *optional*): + Position of the CLS token for each sentence in the batch. If `None`, takes the last token. + is_impossible (`torch.LongTensor` of shape `(batch_size,)`, *optional*): + Whether the question has a possible answer in the paragraph or not. + p_mask (`torch.FloatTensor` of shape `(batch_size, seq_len)`, *optional*): + Mask for tokens at invalid position, such as query and special symbols (PAD, SEP, CLS). 1.0 means token + should be masked. + return_dict (`bool`, *optional*, defaults to `False`): + Whether or not to return a [`~utils.ModelOutput`] instead of a plain tuple. + + Returns: + """ + start_logits = self.start_logits(hidden_states, p_mask=p_mask) + + if start_positions is not None and end_positions is not None: + # If we are on multi-GPU, let's remove the dimension added by batch splitting + for x in (start_positions, end_positions, cls_index, is_impossible): + if x is not None and x.dim() > 1: + x.squeeze_(-1) + + # during training, compute the end logits based on the ground truth of the start position + end_logits = self.end_logits(hidden_states, start_positions=start_positions, p_mask=p_mask) + + loss_fct = CrossEntropyLoss() + start_loss = loss_fct(start_logits, start_positions) + end_loss = loss_fct(end_logits, end_positions) + total_loss = (start_loss + end_loss) / 2 + + if cls_index is not None and is_impossible is not None: + # Predict answerability from the representation of CLS and START + cls_logits = self.answer_class(hidden_states, start_positions=start_positions, cls_index=cls_index) + loss_fct_cls = nn.BCEWithLogitsLoss() + cls_loss = loss_fct_cls(cls_logits, is_impossible) + + # note(zhiliny): by default multiply the loss by 0.5 so that the scale is comparable to start_loss and end_loss + total_loss += cls_loss * 0.5 + + return SquadHeadOutput(loss=total_loss) if return_dict else (total_loss,) + + else: + # during inference, compute the end logits based on beam search + bsz, slen, hsz = hidden_states.size() + start_log_probs = nn.functional.softmax(start_logits, dim=-1) # shape (bsz, slen) + + start_top_log_probs, start_top_index = torch.topk( + start_log_probs, self.start_n_top, dim=-1 + ) # shape (bsz, start_n_top) + start_top_index_exp = start_top_index.unsqueeze(-1).expand(-1, -1, hsz) # shape (bsz, start_n_top, hsz) + start_states = torch.gather(hidden_states, -2, start_top_index_exp) # shape (bsz, start_n_top, hsz) + start_states = start_states.unsqueeze(1).expand(-1, slen, -1, -1) # shape (bsz, slen, start_n_top, hsz) + + hidden_states_expanded = hidden_states.unsqueeze(2).expand_as( + start_states + ) # shape (bsz, slen, start_n_top, hsz) + p_mask = p_mask.unsqueeze(-1) if p_mask is not None else None + end_logits = self.end_logits(hidden_states_expanded, start_states=start_states, p_mask=p_mask) + end_log_probs = nn.functional.softmax(end_logits, dim=1) # shape (bsz, slen, start_n_top) + + end_top_log_probs, end_top_index = torch.topk( + end_log_probs, self.end_n_top, dim=1 + ) # shape (bsz, end_n_top, start_n_top) + end_top_log_probs = end_top_log_probs.view(-1, self.start_n_top * self.end_n_top) + end_top_index = end_top_index.view(-1, self.start_n_top * self.end_n_top) + + start_states = torch.einsum("blh,bl->bh", hidden_states, start_log_probs) + cls_logits = self.answer_class(hidden_states, start_states=start_states, cls_index=cls_index) + + if not return_dict: + return (start_top_log_probs, start_top_index, end_top_log_probs, end_top_index, cls_logits) + else: + return SquadHeadOutput( + start_top_log_probs=start_top_log_probs, + start_top_index=start_top_index, + end_top_log_probs=end_top_log_probs, + end_top_index=end_top_index, + cls_logits=cls_logits, + ) + + +class SequenceSummary(nn.Module): + r""" + Compute a single vector summary of a sequence hidden states. + + Args: + config ([`PretrainedConfig`]): + The config used by the model. Relevant arguments in the config class of the model are (refer to the actual + config class of your model for the default values it uses): + + - **summary_type** (`str`) -- The method to use to make this summary. Accepted values are: + + - `"last"` -- Take the last token hidden state (like XLNet) + - `"first"` -- Take the first token hidden state (like Bert) + - `"mean"` -- Take the mean of all tokens hidden states + - `"cls_index"` -- Supply a Tensor of classification token position (GPT/GPT-2) + - `"attn"` -- Not implemented now, use multi-head attention + + - **summary_use_proj** (`bool`) -- Add a projection after the vector extraction. + - **summary_proj_to_labels** (`bool`) -- If `True`, the projection outputs to `config.num_labels` classes + (otherwise to `config.hidden_size`). + - **summary_activation** (`Optional[str]`) -- Set to `"tanh"` to add a tanh activation to the output, + another string or `None` will add no activation. + - **summary_first_dropout** (`float`) -- Optional dropout probability before the projection and activation. + - **summary_last_dropout** (`float`)-- Optional dropout probability after the projection and activation. + """ + + def __init__(self, config: PretrainedConfig): + super().__init__() + + self.summary_type = getattr(config, "summary_type", "last") + if self.summary_type == "attn": + # We should use a standard multi-head attention module with absolute positional embedding for that. + # Cf. https://github.com/zihangdai/xlnet/blob/master/modeling.py#L253-L276 + # We can probably just use the multi-head attention module of PyTorch >=1.1.0 + raise NotImplementedError + + self.summary = Identity() + if hasattr(config, "summary_use_proj") and config.summary_use_proj: + if hasattr(config, "summary_proj_to_labels") and config.summary_proj_to_labels and config.num_labels > 0: + num_classes = config.num_labels + else: + num_classes = config.hidden_size + self.summary = nn.Linear(config.hidden_size, num_classes) + + activation_string = getattr(config, "summary_activation", None) + self.activation: Callable = get_activation(activation_string) if activation_string else Identity() + + self.first_dropout = Identity() + if hasattr(config, "summary_first_dropout") and config.summary_first_dropout > 0: + self.first_dropout = nn.Dropout(config.summary_first_dropout) + + self.last_dropout = Identity() + if hasattr(config, "summary_last_dropout") and config.summary_last_dropout > 0: + self.last_dropout = nn.Dropout(config.summary_last_dropout) + + def forward( + self, hidden_states: torch.FloatTensor, cls_index: Optional[torch.LongTensor] = None + ) -> torch.FloatTensor: + """ + Compute a single vector summary of a sequence hidden states. + + Args: + hidden_states (`torch.FloatTensor` of shape `[batch_size, seq_len, hidden_size]`): + The hidden states of the last layer. + cls_index (`torch.LongTensor` of shape `[batch_size]` or `[batch_size, ...]` where ... are optional leading dimensions of `hidden_states`, *optional*): + Used if `summary_type == "cls_index"` and takes the last token of the sequence as classification token. + + Returns: + `torch.FloatTensor`: The summary of the sequence hidden states. + """ + if self.summary_type == "last": + output = hidden_states[:, -1] + elif self.summary_type == "first": + output = hidden_states[:, 0] + elif self.summary_type == "mean": + output = hidden_states.mean(dim=1) + elif self.summary_type == "cls_index": + if cls_index is None: + cls_index = torch.full_like( + hidden_states[..., :1, :], + hidden_states.shape[-2] - 1, + dtype=torch.long, + ) + else: + cls_index = cls_index.unsqueeze(-1).unsqueeze(-1) + cls_index = cls_index.expand((-1,) * (cls_index.dim() - 1) + (hidden_states.size(-1),)) + # shape of cls_index: (bsz, XX, 1, hidden_size) where XX are optional leading dim of hidden_states + output = hidden_states.gather(-2, cls_index).squeeze(-2) # shape (bsz, XX, hidden_size) + elif self.summary_type == "attn": + raise NotImplementedError + + output = self.first_dropout(output) + output = self.summary(output) + output = self.activation(output) + output = self.last_dropout(output) + + return output + + +def unwrap_model(model: nn.Module) -> nn.Module: + """ + Recursively unwraps a model from potential containers (as used in distributed training). + + Args: + model (`torch.nn.Module`): The model to unwrap. + """ + # since there could be multiple levels of wrapping, unwrap recursively + if hasattr(model, "module"): + return unwrap_model(model.module) + else: + return model + + +def expand_device_map(device_map, param_names): + """ + Expand a device map to return the correspondance parameter name to device. + """ + new_device_map = {} + for module, device in device_map.items(): + new_device_map.update({p: device for p in param_names if p == module or p.startswith(f"{module}.")}) + return new_device_map + + +def get_disk_only_shard_files(device_map, sharded_metadata): + """ + Returns the list of shard files containing only weights offloaded to disk. + """ + files_content = collections.defaultdict(list) + for weight_name, filename in sharded_metadata["weight_map"].items(): + while len(weight_name) > 0 and weight_name not in device_map: + weight_name = ".".join(weight_name.split(".")[:-1]) + files_content[filename].append(device_map[weight_name]) + + return [fname for fname, devices in files_content.items() if set(devices) == {"disk"}] \ No newline at end of file diff --git a/PyTorch/built-in/foundation/CodeGeeX2/model/modeling_chatglm.py b/PyTorch/built-in/foundation/CodeGeeX2/model/modeling_chatglm.py new file mode 100644 index 0000000000000000000000000000000000000000..fbee00ed85a91ffc5f94ff93e47c27155fe99fc8 --- /dev/null +++ b/PyTorch/built-in/foundation/CodeGeeX2/model/modeling_chatglm.py @@ -0,0 +1,1306 @@ +# coding=utf-8 +# Copyright 2023 Huawei Technologies Co., Ltd +# +# 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. + +""" PyTorch ChatGLM model. """ + +import math +import copy +import os +import warnings +import re +import sys + +import torch +import torch_npu +import torch.utils.checkpoint +import torch.nn.functional as F +from torch import nn +from torch.nn import CrossEntropyLoss, LayerNorm +from torch.nn.utils import skip_init +from typing import Optional, Tuple, Union, List, Callable, Dict, Any + +from transformers.modeling_outputs import ( + BaseModelOutputWithPast, + CausalLMOutputWithPast, +) +from transformers.modeling_utils import PreTrainedModel +from transformers.utils import logging +from transformers.generation.logits_process import LogitsProcessor +from transformers.generation.utils import LogitsProcessorList, StoppingCriteriaList, GenerationConfig, ModelOutput + +from .configuration_chatglm import ChatGLMConfig +from deepspeed.runtime.activation_checkpointing import checkpointing + +# flags required to enable jit fusion kernels +USE_OPTI_TEST = False +USE_SCALED_SOFTMAX = False +USE_FLASH = True +TEST_PREFIX = False +USE_NPU_RMSNORM=True +if sys.platform != 'darwin': + torch._C._jit_set_profiling_mode(False) + torch._C._jit_set_profiling_executor(False) + torch._C._jit_override_can_fuse_on_cpu(True) + torch._C._jit_override_can_fuse_on_gpu(True) + +logger = logging.get_logger(__name__) + +_CHECKPOINT_FOR_DOC = "THUDM/ChatGLM2-6B" +_CONFIG_FOR_DOC = "ChatGLM6BConfig" + +CHATGLM_6B_PRETRAINED_MODEL_ARCHIVE_LIST = [ + "THUDM/chatglm2-6b", + # See all ChatGLM models at https://huggingface.co/models?filter=chatglm +] + + +def default_init(cls, *args, **kwargs): + return cls(*args, **kwargs) + + +class InvalidScoreLogitsProcessor(LogitsProcessor): + def __call__(self, input_ids: torch.LongTensor, scores: torch.FloatTensor) -> torch.FloatTensor: + if torch.isnan(scores).any() or torch.isinf(scores).any(): + scores.zero_() + scores[..., 5] = 5e4 + return scores + + +class PrefixEncoder(torch.nn.Module): + """ + The torch.nn model to encode the prefix + Input shape: (batch-size, prefix-length) + Output shape: (batch-size, prefix-length, 2*layers*hidden) + """ + + def __init__(self, config: ChatGLMConfig): + super().__init__() + self.prefix_projection = config.prefix_projection + if self.prefix_projection: + # Use a two-layer MLP to encode the prefix + kv_size = config.num_layers * config.kv_channels * config.multi_query_group_num * 2 + self.embedding = torch.nn.Embedding(config.pre_seq_len, kv_size) + self.trans = torch.nn.Sequential( + torch.nn.Linear(kv_size, config.hidden_size), + torch.nn.Tanh(), + torch.nn.Linear(config.hidden_size, kv_size) + ) + else: + self.embedding = torch.nn.Embedding(config.pre_seq_len, + config.num_layers * config.kv_channels * config.multi_query_group_num * 2) + + def forward(self, prefix: torch.Tensor): + if self.prefix_projection: + prefix_tokens = self.embedding(prefix) + past_key_values = self.trans(prefix_tokens) + else: + past_key_values = self.embedding(prefix) + return past_key_values + + +def split_tensor_along_last_dim( + tensor: torch.Tensor, + num_partitions: int, + contiguous_split_chunks: bool = False, +) -> List[torch.Tensor]: + """Split a tensor along its last dimension. + + Arguments: + tensor: input tensor. + num_partitions: number of partitions to split the tensor + contiguous_split_chunks: If True, make each chunk contiguous + in memory. + + Returns: + A list of Tensors + """ + # Get the size and dimension. + last_dim = tensor.dim() - 1 + last_dim_size = tensor.size()[last_dim] // num_partitions + # Split. + tensor_list = torch.split(tensor, last_dim_size, dim=last_dim) + # Note: torch.split does not create contiguous tensors by default. + if contiguous_split_chunks: + return tuple(chunk.contiguous() for chunk in tensor_list) + + return tensor_list + + +class RotaryEmbedding(nn.Module): + def __init__(self, dim, original_impl=False, device=None, dtype=None): + super().__init__() + inv_freq = 1.0 / (10000 ** (torch.arange(0, dim, 2, device=device).to(dtype=dtype) / dim).double()).to( + dtype=dtype) + + self.register_buffer("inv_freq", inv_freq) + self.dim = dim + self.original_impl = original_impl + + def forward_impl( + self, seq_len: int, n_elem: int, dtype: torch.dtype, device: torch.device, base: int = 10000 + ): + """Enhanced Transformer with Rotary Position Embedding. + + Derived from: https://github.com/labmlai/annotated_deep_learning_paper_implementations/blob/master/labml_nn/ + transformers/rope/__init__.py. MIT License: + https://github.com/labmlai/annotated_deep_learning_paper_implementations/blob/master/license. + """ + # $\Theta = {\theta_i = 10000^{\frac{2(i-1)}{d}}, i \in [1, 2, ..., \frac{d}{2}]}$ + theta = 1.0 / (base ** (torch.arange(0, n_elem, 2, dtype=dtype, device=device) / n_elem).double()).to( + dtype=dtype) + + # Create position indexes `[0, 1, ..., seq_len - 1]` + seq_idx = torch.arange(seq_len, dtype=dtype, device=device) + + # Calculate the product of position index and $\theta_i$ + idx_theta = torch.outer(seq_idx, theta).float() + cache = torch.stack([torch.cos(idx_theta), torch.sin(idx_theta)], dim=-1) + + # this is to mimic the behaviour of complex32, else we will get different results + if dtype in (torch.float16, torch.bfloat16, torch.int8): + cache = cache.bfloat16() if dtype == torch.bfloat16 else cache.half() + return cache + + def forward(self, max_seq_len, offset=0): + return self.forward_impl( + max_seq_len, self.dim, dtype=self.inv_freq.dtype, device=self.inv_freq.device + ) + + +@torch.jit.script +def apply_rotary_pos_emb(x: torch.Tensor, rope_cache: torch.Tensor) -> torch.Tensor: + # x: [sq, b, np, hn] + # Tag：q, k: [sq, b, np, hn], hn=128 + sq, b, np, hn = x.size(0), x.size(1), x.size(2), x.size(3) + + if USE_OPTI_TEST: + (cos, sin, rot_dim) = rope_cache + x, x_pass = torch.chunk(x, 2, dim=-1) + # rot_dim =64 + # rot_dim = rope_cache.shape[-2] * 2 + # Tag: x-->(sq, b, 64, 2) + # x, x_pass = x[..., :rot_dim], x[..., rot_dim:] + # truncate to support variable sizes + # rope_cache = rope_cache[:sq] + + # (sq, b, 32, 32, 2) + xshaped = x.reshape(sq, -1, np, rot_dim // 2, 2) + # rope_cache = rope_cache.view(sq, -1, 1, xshaped.size(3), 2) + # rope_cache[..., 0]; shape-->(sq, b, 32) + x_out2 = torch.stack( + [ + # [q1, q3, ] * [cos1, cos2, cos3] - [q2, q4, ] * [sin1, sin2, sin3] + xshaped[..., 0] * cos - xshaped[..., 1] * sin, + # [q2, q4, ] * [cos1, cos2, cos3] - [q1, q3, ] * [sin1, sin2, sin3] + xshaped[..., 1] * cos + xshaped[..., 0] * sin, + ], + -1, + ) + else: + # rot_dim =64 + + rot_dim = rope_cache.shape[-2] * 2 + # Tag: x-->(sq, b, 64, 2) + x, x_pass = x[..., :rot_dim], x[..., rot_dim:] + # truncate to support variable sizes + rope_cache = rope_cache[:sq] + + # (sq, b, 32, 32, 2); + xshaped = x.reshape(sq, -1, np, rot_dim // 2, 2) + rope_cache = rope_cache.view(sq, -1, 1, xshaped.size(3), 2) + # rope_cache[..., 0]; shape-->(sq, b, 32) + x_out2 = torch.stack( + [ + # [q1, q3, ] * [cos1, cos2, cos3] - [q2, q4, ] * [sin1, sin2, sin3] + # "2048,1,32,32; 2048,1,1,32" + xshaped[..., 0] * rope_cache[..., 0] - xshaped[..., 1] * rope_cache[..., 1], + # [q2, q4, ] * [cos1, cos2, cos3] - [q1, q3, ] * [sin1, sin2, sin3] + xshaped[..., 1] * rope_cache[..., 0] + xshaped[..., 0] * rope_cache[..., 1], + ], + -1, + ) + + x_out2 = x_out2.flatten(3) + return torch.cat((x_out2, x_pass), dim=-1) + + +class RMSNorm(torch.nn.Module): + def __init__(self, normalized_shape, eps=1e-5, device=None, dtype=None, **kwargs): + super().__init__() + self.weight = torch.nn.Parameter(torch.empty(normalized_shape, device=device, dtype=dtype)) + self.eps = eps + + def forward(self, hidden_states: torch.Tensor): + if USE_NPU_RMSNORM: + return torch.npu_rms_norm(hidden_states, self.weight)[0] + else: + input_dtype = hidden_states.dtype + hidden_states_fp32 = hidden_states.to(torch.float32) + variance = torch.mul(hidden_states_fp32, hidden_states_fp32).mean(-1, keepdim=True) + hidden_states = hidden_states * torch.rsqrt(variance + self.eps).half() + + return (self.weight * hidden_states).to(input_dtype) + + + + +class CoreAttention(torch.nn.Module): + def __init__(self, config: ChatGLMConfig, layer_number): + super(CoreAttention, self).__init__() + + self.apply_query_key_layer_scaling = config.apply_query_key_layer_scaling + self.attention_softmax_in_fp32 = config.attention_softmax_in_fp32 + if self.apply_query_key_layer_scaling: + self.attention_softmax_in_fp32 = True + self.layer_number = max(1, layer_number) + + projection_size = config.kv_channels * config.num_attention_heads + + # Per attention head and per partition values. + self.hidden_size_per_partition = projection_size + self.hidden_size_per_attention_head = projection_size // config.num_attention_heads + self.num_attention_heads_per_partition = config.num_attention_heads + + coeff = None + self.norm_factor = math.sqrt(self.hidden_size_per_attention_head) + if self.apply_query_key_layer_scaling: + coeff = self.layer_number + self.norm_factor *= coeff + self.coeff = coeff + + self.attention_dropout = torch.nn.Dropout(config.attention_dropout) + + def forward(self, query_layer, key_layer, value_layer, attention_mask): + pytorch_major_version = int(torch.__version__.split('.')[0]) + if pytorch_major_version >= 2: + query_layer, key_layer, value_layer = [k.permute(1, 2, 0, 3) for k in [query_layer, key_layer, value_layer]] + if attention_mask is None and query_layer.shape[2] == key_layer.shape[2]: + context_layer = torch.nn.functional.scaled_dot_product_attention(query_layer, key_layer, value_layer, + is_causal=True) + else: + if attention_mask is not None: + attention_mask = ~attention_mask + context_layer = torch.nn.functional.scaled_dot_product_attention(query_layer, key_layer, value_layer, + attention_mask) + context_layer = context_layer.permute(2, 0, 1, 3) + new_context_layer_shape = context_layer.size()[:-2] + (self.hidden_size_per_partition,) + context_layer = context_layer.reshape(*new_context_layer_shape) + else: + # Raw attention scores + # [b, np, sq, sk] + output_size = (query_layer.size(1), query_layer.size(2), query_layer.size(0), key_layer.size(0)) + + alpha = (1.0 / self.norm_factor) + query_layer = query_layer * alpha + if USE_FLASH: + # Tag:-1表示自动计算该维度的大小, 合并维度 + query_layer = query_layer.view(output_size[2], output_size[0], -1) + # (8192, 1, 4096 + key_layer = key_layer.view(output_size[2], output_size[0], -1) + value_layer = value_layer.view(output_size[2], output_size[0], -1) + input_layout = "SBH" + scale = self.coeff + keep_prob = 1.0 + N = output_size[1] + + if TEST_PREFIX == False: + out = torch_npu.npu_flash_attention( \ + query_layer, key_layer, value_layer, N, \ + pse=None, \ + padding_mask=None, \ + atten_mask=attention_mask, \ + scale=scale, \ + keep_prob=keep_prob, \ + input_layout=input_layout, \ + pre_tockens=65536, \ + next_tockens=0)[0] + else: + sparse_mode = 5 + # seq_lenth + past_length = attention_mask.size(3) - query_layer.size(0) + prefixList = [] + # batch + for i in range(0, output_size[0]): + prefixList.append(past_length) + prefix = tuple(prefixList) + out = torch_npu.npu_fusion_attention( + query_layer, key_layer, value_layer, N, \ + pse=None, + padding_mask=None, + atten_mask=attention_mask, + scale=scale, + keep_prob=keep_prob, + input_layout=input_layout, + pre_tockens=65536, + next_tockens=0, + inner_precise=0, + prefix=prefix, + sparse_mode=sparse_mode)[0] + context_layer = out.view(output_size[2], output_size[0], output_size[1], -1) + else: + query_layer = query_layer.view(output_size[2], output_size[0] * output_size[1], -1) + # [sk, b, np, hn] -> [sk, b * np, hn] + key_layer = key_layer.view(output_size[3], output_size[0] * output_size[1], -1) + # preallocting input tensor: [b * np, sq, sk] + matmul_result = torch.bmm(query_layer.transpose(0, 1), key_layer.permute(1, 2, 0)) + # change view to [b, np, sq, sk] + attention_scores = matmul_result.view(*output_size) + # =========================== + # Attention probs and dropout + # =========================== + # attention scores and attention mask [b, np, sq, sk] + + if USE_SCALED_SOFTMAX: + assert self.coeff, attention_mask + attention_probs = torch_npu.npu_scaled_masked_softmax(attention_scores, attention_mask, self.coeff, + False) + else: + if attention_mask is not None: + attention_scores = attention_scores.masked_fill(attention_mask, float("-inf")) + if self.coeff is not None: + attention_scores = attention_scores * self.coeff + attention_probs = F.softmax(attention_scores, dim=-1) + attention_probs = attention_probs.type_as(value_layer) + # This is actually dropping out entire tokens to attend to, which might + # seem a bit unusual, but is taken from the original Transformer paper. + attention_probs = self.attention_dropout(attention_probs) + # ========================= + # Context layer. [sq, b, hp] + # ========================= + # value_layer -> context layer. + # [sk, b, np, hn] --> [b, np, sq, hn] + # context layer shape: [b, np, sq, hn] + output_size = (value_layer.size(1), value_layer.size(2), query_layer.size(0), value_layer.size(3)) + # change view [sk, b * np, hn] + value_layer = value_layer.view(value_layer.size(0), output_size[0] * output_size[1], -1) + # change view [b * np, sq, sk] + attention_probs = attention_probs.view(output_size[0] * output_size[1], output_size[2], -1) + # matmul: [b * np, sq, hn] + context_layer = torch.bmm(attention_probs, value_layer.transpose(0, 1)) + # change view [b, np, sq, hn] + context_layer = context_layer.view(*output_size) + # [b, np, sq, hn] --> [sq, b, np, hn] + # context_layer = context_layer.permute(2, 0, 1, 3).contiguous() + context_layer = torch.npu_confusion_transpose(context_layer, [2, 0, 1, 3], [*output_size], False) + # [sq, b, np, hn] --> [sq, b, hp] + new_context_layer_shape = context_layer.size()[:-2] + (self.hidden_size_per_partition,) + context_layer = context_layer.view(*new_context_layer_shape) + return context_layer + + +class SelfAttention(torch.nn.Module): + """Parallel self-attention layer abstract class. + + Self-attention layer takes input with size [s, b, h] + and returns output of the same size. + """ + + def __init__(self, config: ChatGLMConfig, layer_number, device=None): + super(SelfAttention, self).__init__() + self.layer_number = max(1, layer_number) + self.projection_size = config.kv_channels * config.num_attention_heads + # Per attention head and per partition values. + # 128 + self.hidden_size_per_attention_head = self.projection_size // config.num_attention_heads + self.num_attention_heads_per_partition = config.num_attention_heads + self.multi_query_attention = config.multi_query_attention + self.qkv_hidden_size = 3 * self.projection_size + if self.multi_query_attention: + self.num_multi_query_groups_per_partition = config.multi_query_group_num + self.qkv_hidden_size = ( + self.projection_size + 2 * self.hidden_size_per_attention_head * config.multi_query_group_num + ) + self.query_key_value = nn.Linear(config.hidden_size, self.qkv_hidden_size, + bias=config.add_bias_linear or config.add_qkv_bias, + device=device, **_config_to_kwargs(config) + ) + + self.core_attention = CoreAttention(config, self.layer_number) + + # Output. + self.dense = nn.Linear(self.projection_size, config.hidden_size, bias=config.add_bias_linear, + device=device, **_config_to_kwargs(config) + ) + + def _allocate_memory(self, inference_max_sequence_len, batch_size, device=None, dtype=None): + if self.multi_query_attention: + num_attention_heads = self.num_multi_query_groups_per_partition + else: + num_attention_heads = self.num_attention_heads_per_partition + return torch.empty( + inference_max_sequence_len, + batch_size, + num_attention_heads, + self.hidden_size_per_attention_head, + dtype=dtype, + device=device, + ) + + def forward( + self, hidden_states, attention_mask, rotary_pos_emb, kv_cache=None, use_cache=True + ): + # hidden_states: [sq, b, h] + + # ================================================= + # Pre-allocate memory for key-values for inference. + # ================================================= + # ===================== + # Query, Key, and Value + # ===================== + + # Attention heads [sq, b, h] --> [sq, b, (np * 3 * hn)] + # qvk_hidden_size = 4096 +2*128*2 + # self.query_key_value = nn.Linear(config.hidden_size, self.qkv_hidden_size, + # bias=config.add_bias_linear or config.add_qkv_bias, + # device=device, **_config_to_kwargs(config) + # ) + # 4096, 4, 4096 -->4096, 4, 4608 + mixed_x_layer = self.query_key_value(hidden_states) + if self.multi_query_attention: + (query_layer, key_layer, value_layer) = mixed_x_layer.split( + [ + self.num_attention_heads_per_partition * self.hidden_size_per_attention_head, + self.num_multi_query_groups_per_partition * self.hidden_size_per_attention_head, + self.num_multi_query_groups_per_partition * self.hidden_size_per_attention_head, + ], + dim=-1, + ) + query_layer = query_layer.view( + query_layer.size()[:-1] + (self.num_attention_heads_per_partition, self.hidden_size_per_attention_head) + ) + key_layer = key_layer.view( + key_layer.size()[:-1] + (self.num_multi_query_groups_per_partition, self.hidden_size_per_attention_head) + ) + value_layer = value_layer.view( + value_layer.size()[:-1] + + (self.num_multi_query_groups_per_partition, self.hidden_size_per_attention_head) + ) + else: + new_tensor_shape = mixed_x_layer.size()[:-1] + \ + (self.num_attention_heads_per_partition, + 3 * self.hidden_size_per_attention_head) + mixed_x_layer = mixed_x_layer.view(*new_tensor_shape) + + # [sq, b, np, 3 * hn] --> 3 [sq, b, np, hn] + (query_layer, key_layer, value_layer) = split_tensor_along_last_dim(mixed_x_layer, 3) + + # apply relative positional encoding (rotary embedding) + if rotary_pos_emb is not None: + query_layer = apply_rotary_pos_emb(query_layer, rotary_pos_emb) + key_layer = apply_rotary_pos_emb(key_layer, rotary_pos_emb) + + # adjust key and value for inference + if kv_cache is not None: + cache_k, cache_v = kv_cache + key_layer = torch.cat((cache_k, key_layer), dim=0) + value_layer = torch.cat((cache_v, value_layer), dim=0) + if use_cache: + kv_cache = (key_layer, value_layer) + else: + kv_cache = None + + if self.multi_query_attention: + key_layer = key_layer.unsqueeze(-2) + key_layer = key_layer.expand( + -1, -1, -1, self.num_attention_heads_per_partition // self.num_multi_query_groups_per_partition, -1 + ) + key_layer = key_layer.contiguous().view( + key_layer.size()[:2] + (self.num_attention_heads_per_partition, self.hidden_size_per_attention_head) + ) + value_layer = value_layer.unsqueeze(-2) + value_layer = value_layer.expand( + -1, -1, -1, self.num_attention_heads_per_partition // self.num_multi_query_groups_per_partition, -1 + ) + value_layer = value_layer.contiguous().view( + value_layer.size()[:2] + (self.num_attention_heads_per_partition, self.hidden_size_per_attention_head) + ) + + # ================================== + # core attention computation + # ================================== + + context_layer = self.core_attention(query_layer, key_layer, value_layer, attention_mask) + + # ================= + # Output. [sq, b, h] + # ================= + + output = self.dense(context_layer) + + return output, kv_cache + + +def _config_to_kwargs(args): + common_kwargs = { + "dtype": args.torch_dtype, + } + return common_kwargs + + +class MLP(torch.nn.Module): + """MLP. + + MLP will take the input with h hidden state, project it to 4*h + hidden dimension, perform nonlinear transformation, and project the + state back into h hidden dimension. + """ + + def __init__(self, config: ChatGLMConfig, device=None): + super(MLP, self).__init__() + + self.add_bias = config.add_bias_linear + + # Project to 4h. If using swiglu double the output width, see https://arxiv.org/pdf/2002.05202.pdf + self.dense_h_to_4h = nn.Linear( + config.hidden_size, + config.ffn_hidden_size * 2, + bias=self.add_bias, + device=device, + **_config_to_kwargs(config) + ) + + def swiglu(x): + x = torch.chunk(x, 2, dim=-1) + return F.silu(x[0]) * x[1] + + self.activation_func = swiglu + + # Project back to h. + self.dense_4h_to_h = nn.Linear( + config.ffn_hidden_size, + config.hidden_size, + bias=self.add_bias, + device=device, + **_config_to_kwargs(config) + ) + + def forward(self, hidden_states): + # [s, b, 4hp] + intermediate_parallel = self.dense_h_to_4h(hidden_states) + intermediate_parallel = self.activation_func(intermediate_parallel) + # [s, b, h] + output = self.dense_4h_to_h(intermediate_parallel) + return output + + +class GLMBlock(torch.nn.Module): + """A single transformer layer. + + Transformer layer takes input with size [s, b, h] and returns an + output of the same size. + """ + + def __init__(self, config: ChatGLMConfig, layer_number, device=None): + super(GLMBlock, self).__init__() + self.layer_number = layer_number + + self.apply_residual_connection_post_layernorm = config.apply_residual_connection_post_layernorm + + self.fp32_residual_connection = config.fp32_residual_connection + + LayerNormFunc = RMSNorm if config.rmsnorm else LayerNorm + # Layernorm on the input data. + self.input_layernorm = LayerNormFunc(config.hidden_size, eps=config.layernorm_epsilon, device=device, + dtype=config.torch_dtype) + + # Self attention. + self.self_attention = SelfAttention(config, layer_number, device=device) + self.hidden_dropout = config.hidden_dropout + + # Layernorm on the attention output + self.post_attention_layernorm = LayerNormFunc(config.hidden_size, eps=config.layernorm_epsilon, device=device, + dtype=config.torch_dtype) + + # MLP + self.mlp = MLP(config, device=device) + + def forward( + self, hidden_states, attention_mask, rotary_pos_emb, kv_cache=None, use_cache=True, + ): + # hidden_states: [s, b, h] + # Layer norm at the beginning of the transformer layer. + layernorm_output = self.input_layernorm(hidden_states) + # Self attention. + attention_output, kv_cache = self.self_attention( + layernorm_output, + attention_mask, + rotary_pos_emb, + kv_cache=kv_cache, + use_cache=use_cache + ) + + # Residual connection. + if self.apply_residual_connection_post_layernorm: + residual = layernorm_output + else: + residual = hidden_states + + layernorm_input = torch.nn.functional.dropout(attention_output, p=self.hidden_dropout, training=self.training) + layernorm_input = residual + layernorm_input + + # Layer norm post the self attention. + layernorm_output = self.post_attention_layernorm(layernorm_input) + + # MLP. + mlp_output = self.mlp(layernorm_output) + + # Second residual connection. + if self.apply_residual_connection_post_layernorm: + residual = layernorm_output + else: + residual = layernorm_input + + output = torch.nn.functional.dropout(mlp_output, p=self.hidden_dropout, training=self.training) + output = residual + output + + return output, kv_cache + + +class GLMTransformer(torch.nn.Module): + """Transformer class.""" + + def __init__(self, config: ChatGLMConfig, device=None): + super(GLMTransformer, self).__init__() + + self.fp32_residual_connection = config.fp32_residual_connection + self.post_layer_norm = config.post_layer_norm + + # Number of layers. + self.num_layers = config.num_layers + + # Transformer layers. + def build_layer(layer_number): + return GLMBlock(config, layer_number, device=device) + + self.layers = torch.nn.ModuleList([build_layer(i + 1) for i in range(self.num_layers)]) + + if self.post_layer_norm: + LayerNormFunc = RMSNorm if config.rmsnorm else LayerNorm + # Final layer norm before output. + self.final_layernorm = LayerNormFunc(config.hidden_size, eps=config.layernorm_epsilon, device=device, + dtype=config.torch_dtype) + + self.gradient_checkpointing = False + + def _get_layer(self, layer_number): + return self.layers[layer_number] + + def forward( + self, hidden_states, attention_mask, rotary_pos_emb, kv_caches=None, + use_cache: Optional[bool] = True, + output_hidden_states: Optional[bool] = False, + ): + if not kv_caches: + kv_caches = [None for _ in range(self.num_layers)] + presents = () if use_cache else None + if self.gradient_checkpointing and self.training: + if use_cache: + logger.warning_once( + "`use_cache=True` is incompatible with gradient checkpointing. Setting `use_cache=False`..." + ) + use_cache = False + + all_self_attentions = None + all_hidden_states = () if output_hidden_states else None + for index in range(self.num_layers): + if output_hidden_states: + all_hidden_states = all_hidden_states + (hidden_states,) + + layer = self._get_layer(index) + if self.gradient_checkpointing and self.training: + layer_ret = checkpointing.checkpoint( + # layer_ret = torch.utils.checkpoint.checkpoint( + layer, + hidden_states, + attention_mask, + rotary_pos_emb, + kv_caches[index], + use_cache + ) + else: + layer_ret = layer( + hidden_states, + attention_mask, + rotary_pos_emb, + kv_cache=kv_caches[index], + use_cache=use_cache + ) + hidden_states, kv_cache = layer_ret + if use_cache: + presents = presents + (kv_cache,) + + if output_hidden_states: + all_hidden_states = all_hidden_states + (hidden_states,) + + # Final layer norm. + if self.post_layer_norm: + hidden_states = self.final_layernorm(hidden_states) + + return hidden_states, presents, all_hidden_states, all_self_attentions + + +class ChatGLMPreTrainedModel(PreTrainedModel): + """ + An abstract class to handle weights initialization and + a simple interface for downloading and loading pretrained models. + """ + + is_parallelizable = False + supports_gradient_checkpointing = True + config_class = ChatGLMConfig + base_model_prefix = "transformer" + _no_split_modules = ["GLMBlock"] + + def _init_weights(self, module: nn.Module): + """Initialize the weights.""" + return + + def get_masks_(self, input_ids, device): + batch_size, seq_length = input_ids.shape + attention_mask = torch.ones(batch_size, 1, seq_length, seq_length, device=input_ids.device) + attention_mask.tril_() + # attention_mask = ~attention_mask + attention_mask = (attention_mask < 0.5).bool() + + return attention_mask + + def get_masks(self, input_ids, past_key_values, padding_mask=None): + batch_size, seq_length = input_ids.shape + full_attention_mask = torch.ones(batch_size, seq_length, seq_length, device=input_ids.device) + full_attention_mask.tril_() + past_length = 0 + if past_key_values: + past_length = past_key_values[0][0].shape[0] + if past_length: + full_attention_mask = torch.cat((torch.ones(batch_size, seq_length, past_length, + device=input_ids.device), full_attention_mask), dim=-1) + if padding_mask is not None: + full_attention_mask = full_attention_mask * padding_mask.unsqueeze(1) + if not past_length and padding_mask is not None: + full_attention_mask -= padding_mask.unsqueeze(-1) - 1 + full_attention_mask = (full_attention_mask < 0.5).bool() + full_attention_mask.unsqueeze_(1) + return full_attention_mask + + def get_position_ids(self, input_ids, device): + batch_size, seq_length = input_ids.shape + position_ids = torch.arange(seq_length, dtype=torch.long, device=device).unsqueeze(0).repeat(batch_size, 1) + return position_ids + + def _set_gradient_checkpointing(self, module, value=False): + if isinstance(module, GLMTransformer): + module.gradient_checkpointing = value + + +class Embedding(torch.nn.Module): + """Language model embeddings.""" + + def __init__(self, config: ChatGLMConfig, device=None): + super(Embedding, self).__init__() + + self.hidden_size = config.hidden_size + # Word embeddings (parallel). + self.word_embeddings = nn.Embedding( + config.padded_vocab_size, + self.hidden_size, + dtype=config.torch_dtype, + device=device + ) + self.fp32_residual_connection = config.fp32_residual_connection + + def forward(self, input_ids): + # Embeddings. + words_embeddings = self.word_embeddings(input_ids) + embeddings = words_embeddings + # Data format change to avoid explicit tranposes : [b s h] --> [s b h]. + embeddings = embeddings.transpose(0, 1).contiguous() + # If the input flag for fp32 residual connection is set, convert for float. + if self.fp32_residual_connection: + embeddings = embeddings.float() + return embeddings + + +class ChatGLMModel(ChatGLMPreTrainedModel): + def __init__(self, config: ChatGLMConfig, device=None, empty_init=True): + super().__init__(config) + if empty_init: + init_method = skip_init + else: + init_method = default_init + init_kwargs = {} + if device is not None: + init_kwargs["device"] = device + self.embedding = init_method(Embedding, config, **init_kwargs) + self.num_layers = config.num_layers + self.multi_query_group_num = config.multi_query_group_num + self.kv_channels = config.kv_channels + + # Rotary positional embeddings + self.seq_length = config.seq_length + rotary_dim = ( + config.hidden_size // config.num_attention_heads if config.kv_channels is None else config.kv_channels + ) + + self.rotary_pos_emb = RotaryEmbedding(rotary_dim // 2, original_impl=config.original_rope, device=device, + dtype=config.torch_dtype) + self.encoder = init_method(GLMTransformer, config, **init_kwargs) + self.output_layer = init_method(nn.Linear, config.hidden_size, config.padded_vocab_size, bias=False, + dtype=config.torch_dtype, **init_kwargs) + self.pre_seq_len = config.pre_seq_len + self.prefix_projection = config.prefix_projection + if self.pre_seq_len is not None: + for param in self.parameters(): + param.requires_grad = False + self.prefix_tokens = torch.arange(self.pre_seq_len).int() + self.prefix_encoder = PrefixEncoder(config) + self.dropout = torch.nn.Dropout(0.1) + + def get_input_embeddings(self): + return self.embedding.word_embeddings + + def get_prompt(self, batch_size, device, dtype=torch.half): + prefix_tokens = self.prefix_tokens.unsqueeze(0).expand(batch_size, -1).to(device) + past_key_values = self.prefix_encoder(prefix_tokens).type(dtype) + past_key_values = past_key_values.view( + batch_size, + self.pre_seq_len, + self.num_layers * 2, + self.multi_query_group_num, + self.kv_channels + ) + # seq_len, b, nh, hidden_size + past_key_values = self.dropout(past_key_values) + past_key_values = past_key_values.permute([2, 1, 0, 3, 4]).split(2) + return past_key_values + + def forward( + self, + input_ids, + position_ids: Optional[torch.Tensor] = None, + attention_mask: Optional[torch.BoolTensor] = None, + full_attention_mask: Optional[torch.BoolTensor] = None, + past_key_values: Optional[Tuple[Tuple[torch.Tensor, torch.Tensor], ...]] = None, + inputs_embeds: Optional[torch.Tensor] = None, + use_cache: Optional[bool] = None, + output_hidden_states: Optional[bool] = None, + return_dict: Optional[bool] = None, + ): + output_hidden_states = ( + output_hidden_states if output_hidden_states is not None else self.config.output_hidden_states + ) + use_cache = use_cache if use_cache is not None else self.config.use_cache + return_dict = return_dict if return_dict is not None else self.config.use_return_dict + + batch_size, seq_length = input_ids.shape + + if inputs_embeds is None: + inputs_embeds = self.embedding(input_ids) + + if self.pre_seq_len is not None: + if past_key_values is None: + past_key_values = self.get_prompt(batch_size=batch_size, device=input_ids.device, + dtype=inputs_embeds.dtype) + if attention_mask is not None: + attention_mask = torch.cat([attention_mask.new_ones((batch_size, self.pre_seq_len)), + attention_mask], dim=-1) + if full_attention_mask is None: + if (attention_mask is not None and not attention_mask.all()) or (past_key_values and seq_length != 1): + full_attention_mask = self.get_masks(input_ids, past_key_values, padding_mask=attention_mask) + else: + full_attention_mask = self.get_masks_(input_ids, device=input_ids.device) + + # Rotary positional embeddings + rotary_pos_emb = self.rotary_pos_emb(self.seq_length) + if position_ids is not None: + rotary_pos_emb = rotary_pos_emb[position_ids] + else: + rotary_pos_emb = rotary_pos_emb[None, :seq_length] + rotary_pos_emb = rotary_pos_emb.transpose(0, 1).contiguous() + if USE_OPTI_TEST: + sq = seq_length + rope_cache = rotary_pos_emb + rot_dim = rope_cache.shape[-2] * 2 + rope_cache = rope_cache.view(sq, -1, 1, rot_dim // 2, 2) + cos = rope_cache[..., 0] + sin = rope_cache[..., 1] + rotary_pos_emb = (cos, sin, rot_dim) + + # Run encoder. + hidden_states, presents, all_hidden_states, all_self_attentions = self.encoder( + inputs_embeds, full_attention_mask, rotary_pos_emb=rotary_pos_emb, + kv_caches=past_key_values, use_cache=use_cache, output_hidden_states=output_hidden_states + ) + + if not return_dict: + return tuple(v for v in [hidden_states, presents, all_hidden_states, all_self_attentions] if v is not None) + + return BaseModelOutputWithPast( + last_hidden_state=hidden_states, + past_key_values=presents, + hidden_states=all_hidden_states, + attentions=all_self_attentions, + ) + + def quantize(self, weight_bit_width: int): + from .quantization import quantize + quantize(self.encoder, weight_bit_width) + return self + + +class ChatGLMForConditionalGeneration(ChatGLMPreTrainedModel): + def __init__(self, config: ChatGLMConfig, empty_init=True, device=None): + super().__init__(config) + + self.max_sequence_length = config.max_length + self.transformer = ChatGLMModel(config, empty_init=empty_init, device=device) + self.config = config + self.quantized = False + + if self.config.quantization_bit: + self.quantize(self.config.quantization_bit, empty_init=True) + + def _update_model_kwargs_for_generation( + self, + outputs: ModelOutput, + model_kwargs: Dict[str, Any], + is_encoder_decoder: bool = False, + standardize_cache_format: bool = False, + ) -> Dict[str, Any]: + # update past_key_values + model_kwargs["past_key_values"] = self._extract_past_from_model_output( + outputs, standardize_cache_format=standardize_cache_format + ) + + # update attention mask + if "attention_mask" in model_kwargs: + attention_mask = model_kwargs["attention_mask"] + model_kwargs["attention_mask"] = torch.cat( + [attention_mask, attention_mask.new_ones((attention_mask.shape[0], 1))], dim=-1 + ) + + # update position ids + if "position_ids" in model_kwargs: + position_ids = model_kwargs["position_ids"] + new_position_id = position_ids[..., -1:].clone() + new_position_id += 1 + model_kwargs["position_ids"] = torch.cat( + [position_ids, new_position_id], dim=-1 + ) + + model_kwargs["is_first_forward"] = False + return model_kwargs + + def prepare_inputs_for_generation( + self, + input_ids: torch.LongTensor, + past_key_values: Optional[torch.Tensor] = None, + attention_mask: Optional[torch.Tensor] = None, + position_ids: Optional[torch.Tensor] = None, + is_first_forward: bool = True, + **kwargs + ) -> dict: + # only last token for input_ids if past is not None + if position_ids is None: + position_ids = self.get_position_ids(input_ids, device=input_ids.device) + if not is_first_forward: + position_ids = position_ids[..., -1:] + input_ids = input_ids[:, -1:] + return { + "input_ids": input_ids, + "past_key_values": past_key_values, + "position_ids": position_ids, + "attention_mask": attention_mask, + "return_last_logit": True + } + + def forward( + self, + input_ids: Optional[torch.Tensor] = None, + position_ids: Optional[torch.Tensor] = None, + attention_mask: Optional[torch.Tensor] = None, + past_key_values: Optional[Tuple[torch.FloatTensor]] = None, + inputs_embeds: Optional[torch.Tensor] = None, + labels: Optional[torch.Tensor] = None, + use_cache: Optional[bool] = None, + output_attentions: Optional[bool] = None, + output_hidden_states: Optional[bool] = None, + return_dict: Optional[bool] = None, + return_last_logit: Optional[bool] = False, + ): + use_cache = use_cache if use_cache is not None else self.config.use_cache + return_dict = return_dict if return_dict is not None else self.config.use_return_dict + + transformer_outputs = self.transformer( + input_ids=input_ids, + position_ids=position_ids, + attention_mask=attention_mask, + past_key_values=past_key_values, + inputs_embeds=inputs_embeds, + use_cache=use_cache, + output_hidden_states=output_hidden_states, + return_dict=return_dict, + ) + + hidden_states = transformer_outputs[0] + if return_last_logit: + hidden_states = hidden_states[-1:] + lm_logits = self.transformer.output_layer(hidden_states) + lm_logits = lm_logits.transpose(0, 1).contiguous() + + loss = None + if labels is not None: + lm_logits = lm_logits.to(torch.float32) + + # Shift so that tokens < n predict n + shift_logits = lm_logits[..., :-1, :].contiguous() + shift_labels = labels[..., 1:].contiguous() + # Flatten the tokens + loss_fct = CrossEntropyLoss(ignore_index=-100) + loss = loss_fct(shift_logits.view(-1, shift_logits.size(-1)), shift_labels.view(-1)) + + lm_logits = lm_logits.to(hidden_states.dtype) + loss = loss.to(hidden_states.dtype) + + if not return_dict: + output = (lm_logits,) + transformer_outputs[1:] + return ((loss,) + output) if loss is not None else output + + return CausalLMOutputWithPast( + loss=loss, + logits=lm_logits, + past_key_values=transformer_outputs.past_key_values, + hidden_states=transformer_outputs.hidden_states, + attentions=transformer_outputs.attentions, + ) + + @staticmethod + def _reorder_cache( + past: Tuple[Tuple[torch.Tensor, torch.Tensor], ...], beam_idx: torch.LongTensor + ) -> Tuple[Tuple[torch.Tensor, torch.Tensor], ...]: + """ + This function is used to re-order the `past_key_values` cache if [`~PreTrainedModel.beam_search`] or + [`~PreTrainedModel.beam_sample`] is called. This is required to match `past_key_values` with the correct + beam_idx at every generation step. + + Output shares the same memory storage as `past`. + """ + return tuple( + ( + layer_past[0].index_select(1, beam_idx.to(layer_past[0].device)), + layer_past[1].index_select(1, beam_idx.to(layer_past[1].device)), + ) + for layer_past in past + ) + + def process_response(self, response): + response = response.strip() + response = response.replace("[[训练时间]]", "2023年") + return response + + def build_inputs(self, tokenizer, query: str, history: List[Tuple[str, str]] = None): + prompt = tokenizer.build_prompt(query, history=history) + inputs = tokenizer([prompt], return_tensors="pt") + inputs = inputs.to(self.device) + return inputs + + def build_stream_inputs(self, tokenizer, query: str, history: List[Tuple[str, str]] = None): + if history: + prompt = "\n\n[Round {}]\n\n问：{}\n\n答：".format(len(history) + 1, query) + input_ids = tokenizer.encode(prompt, add_special_tokens=False) + input_ids = input_ids[1:] + inputs = tokenizer.batch_encode_plus([(input_ids, None)], return_tensors="pt", add_special_tokens=False) + else: + prompt = "[Round {}]\n\n问：{}\n\n答：".format(len(history) + 1, query) + inputs = tokenizer([prompt], return_tensors="pt") + inputs = inputs.to(self.device) + return inputs + + @torch.no_grad() + def chat(self, tokenizer, query: str, history: List[Tuple[str, str]] = None, max_length: int = 8192, num_beams=1, + do_sample=True, top_p=0.8, temperature=0.8, logits_processor=None, **kwargs): + if history is None: + history = [] + if logits_processor is None: + logits_processor = LogitsProcessorList() + logits_processor.append(InvalidScoreLogitsProcessor()) + gen_kwargs = {"max_length": max_length, "num_beams": num_beams, "do_sample": do_sample, "top_p": top_p, + "temperature": temperature, "logits_processor": logits_processor, **kwargs} + inputs = self.build_inputs(tokenizer, query, history=history) + outputs = self.generate(**inputs, **gen_kwargs) + outputs = outputs.tolist()[0][len(inputs["input_ids"][0]):] + response = tokenizer.decode(outputs) + response = self.process_response(response) + history = history + [(query, response)] + return response, history + + @torch.no_grad() + def stream_chat(self, tokenizer, query: str, history: List[Tuple[str, str]] = None, past_key_values=None, + max_length: int = 8192, do_sample=True, top_p=0.8, temperature=0.8, logits_processor=None, + return_past_key_values=False, **kwargs): + if history is None: + history = [] + if logits_processor is None: + logits_processor = LogitsProcessorList() + logits_processor.append(InvalidScoreLogitsProcessor()) + gen_kwargs = {"max_length": max_length, "do_sample": do_sample, "top_p": top_p, + "temperature": temperature, "logits_processor": logits_processor, **kwargs} + if past_key_values is None and not return_past_key_values: + inputs = self.build_inputs(tokenizer, query, history=history) + else: + inputs = self.build_stream_inputs(tokenizer, query, history=history) + if past_key_values is not None: + past_length = past_key_values[0][0].shape[0] + if self.transformer.pre_seq_len is not None: + past_length -= self.transformer.pre_seq_len + inputs.position_ids += past_length + attention_mask = inputs.attention_mask + attention_mask = torch.cat((attention_mask.new_ones(1, past_length), attention_mask), dim=1) + inputs['attention_mask'] = attention_mask + for outputs in self.stream_generate(**inputs, past_key_values=past_key_values, + return_past_key_values=return_past_key_values, **gen_kwargs): + if return_past_key_values: + outputs, past_key_values = outputs + outputs = outputs.tolist()[0][len(inputs["input_ids"][0]):] + response = tokenizer.decode(outputs) + if response and response[-1] != "�": + response = self.process_response(response) + new_history = history + [(query, response)] + if return_past_key_values: + yield response, new_history, past_key_values + else: + yield response, new_history + + @torch.no_grad() + def stream_generate( + self, + input_ids, + generation_config: Optional[GenerationConfig] = None, + logits_processor: Optional[LogitsProcessorList] = None, + stopping_criteria: Optional[StoppingCriteriaList] = None, + prefix_allowed_tokens_fn: Optional[Callable[[int, torch.Tensor], List[int]]] = None, + return_past_key_values=False, + **kwargs, + ): + batch_size, input_ids_seq_length = input_ids.shape[0], input_ids.shape[-1] + + if generation_config is None: + generation_config = self.generation_config + generation_config = copy.deepcopy(generation_config) + model_kwargs = generation_config.update(**kwargs) + bos_token_id, eos_token_id = generation_config.bos_token_id, generation_config.eos_token_id + + if isinstance(eos_token_id, int): + eos_token_id = [eos_token_id] + + has_default_max_length = kwargs.get("max_length") is None and generation_config.max_length is not None + if has_default_max_length and generation_config.max_new_tokens is None: + warnings.warn( + f"Using `max_length`'s default ({generation_config.max_length}) to control the generation length. " + "This behaviour is deprecated and will be removed from the config in v5 of Transformers -- we" + " recommend using `max_new_tokens` to control the maximum length of the generation.", + UserWarning, + ) + elif generation_config.max_new_tokens is not None: + generation_config.max_length = generation_config.max_new_tokens + input_ids_seq_length + if not has_default_max_length: + logger.warn( + f"Both `max_new_tokens` (={generation_config.max_new_tokens}) and `max_length`(=" + f"{generation_config.max_length}) seem to have been set. `max_new_tokens` will take precedence. " + "Please refer to the documentation for more information. " + "(https://huggingface.co/docs/transformers/main/en/main_classes/text_generation)", + UserWarning, + ) + + if input_ids_seq_length >= generation_config.max_length: + input_ids_string = "decoder_input_ids" if self.config.is_encoder_decoder else "input_ids" + logger.warning( + f"Input length of {input_ids_string} is {input_ids_seq_length}, but `max_length` is set to" + f" {generation_config.max_length}. This can lead to unexpected behavior. You should consider" + " increasing `max_new_tokens`." + ) + + # 2. Set generation parameters if not already defined + logits_processor = logits_processor if logits_processor is not None else LogitsProcessorList() + stopping_criteria = stopping_criteria if stopping_criteria is not None else StoppingCriteriaList() + + logits_processor = self._get_logits_processor( + generation_config=generation_config, + input_ids_seq_length=input_ids_seq_length, + encoder_input_ids=input_ids, + prefix_allowed_tokens_fn=prefix_allowed_tokens_fn, + logits_processor=logits_processor, + ) + + stopping_criteria = self._get_stopping_criteria( + generation_config=generation_config, stopping_criteria=stopping_criteria + ) + logits_warper = self._get_logits_warper(generation_config) + + unfinished_sequences = input_ids.new(input_ids.shape[0]).fill_(1) + scores = None + while True: + model_inputs = self.prepare_inputs_for_generation(input_ids, **model_kwargs) + # forward pass to get next token + outputs = self( + **model_inputs, + return_dict=True, + output_attentions=False, + output_hidden_states=False, + ) + + next_token_logits = outputs.logits[:, -1, :] + + # pre-process distribution + next_token_scores = logits_processor(input_ids, next_token_logits) + next_token_scores = logits_warper(input_ids, next_token_scores) + + # sample + probs = nn.functional.softmax(next_token_scores, dim=-1) + if generation_config.do_sample: + next_tokens = torch.multinomial(probs, num_samples=1).squeeze(1) + else: + next_tokens = torch.argmax(probs, dim=-1) + + # update generated ids, model inputs, and length for next step + input_ids = torch.cat([input_ids, next_tokens[:, None]], dim=-1) + model_kwargs = self._update_model_kwargs_for_generation( + outputs, model_kwargs, is_encoder_decoder=self.config.is_encoder_decoder + ) + unfinished_sequences = unfinished_sequences.mul((sum(next_tokens != i for i in eos_token_id)).int()) + if return_past_key_values: + yield input_ids, outputs.past_key_values + else: + yield input_ids + # stop when each sentence is finished, or if we exceed the maximum length + if unfinished_sequences.max() == 0 or stopping_criteria(input_ids, scores): + break + + def quantize(self, bits: int, empty_init=False, device=None, **kwargs): + if bits == 0: + return + + from .quantization import quantize + + if self.quantized: + logger.info("Already quantized.") + return self + + self.quantized = True + + self.config.quantization_bit = bits + + self.transformer.encoder = quantize(self.transformer.encoder, bits, empty_init=empty_init, device=device, + **kwargs) + return self diff --git a/PyTorch/built-in/foundation/CodeGeeX2/openai_api.py b/PyTorch/built-in/foundation/CodeGeeX2/openai_api.py new file mode 100644 index 0000000000000000000000000000000000000000..72255624396b2ac5294fca39e18a7b90c5b8efa9 --- /dev/null +++ b/PyTorch/built-in/foundation/CodeGeeX2/openai_api.py @@ -0,0 +1,177 @@ +# coding=utf-8 +# Implements API for ChatGLM2-6B in OpenAI's format. (https://platform.openai.com/docs/api-reference/chat) +# Usage: python openai_api.py +# Visit http://localhost:8000/docs for documents. + + +import time +import torch +import uvicorn +from pydantic import BaseModel, Field +from fastapi import FastAPI, HTTPException +from fastapi.middleware.cors import CORSMiddleware +from contextlib import asynccontextmanager +from typing import Any, Dict, List, Literal, Optional, Union +from transformers import AutoTokenizer, AutoModel +from sse_starlette.sse import ServerSentEvent, EventSourceResponse + + +@asynccontextmanager +async def lifespan(app: FastAPI): # collects GPU memory + yield + if torch.cuda.is_available(): + torch.cuda.empty_cache() + torch.cuda.ipc_collect() + + +app = FastAPI(lifespan=lifespan) + +app.add_middleware( + CORSMiddleware, + allow_origins=["*"], + allow_credentials=True, + allow_methods=["*"], + allow_headers=["*"], +) + +class ModelCard(BaseModel): + id: str + object: str = "model" + created: int = Field(default_factory=lambda: int(time.time())) + owned_by: str = "owner" + root: Optional[str] = None + parent: Optional[str] = None + permission: Optional[list] = None + + +class ModelList(BaseModel): + object: str = "list" + data: List[ModelCard] = [] + + +class ChatMessage(BaseModel): + role: Literal["user", "assistant", "system"] + content: str + + +class DeltaMessage(BaseModel): + role: Optional[Literal["user", "assistant", "system"]] = None + content: Optional[str] = None + + +class ChatCompletionRequest(BaseModel): + model: str + messages: List[ChatMessage] + temperature: Optional[float] = None + top_p: Optional[float] = None + max_length: Optional[int] = None + stream: Optional[bool] = False + + +class ChatCompletionResponseChoice(BaseModel): + index: int + message: ChatMessage + finish_reason: Literal["stop", "length"] + + +class ChatCompletionResponseStreamChoice(BaseModel): + index: int + delta: DeltaMessage + finish_reason: Optional[Literal["stop", "length"]] + + +class ChatCompletionResponse(BaseModel): + model: str + object: Literal["chat.completion", "chat.completion.chunk"] + choices: List[Union[ChatCompletionResponseChoice, ChatCompletionResponseStreamChoice]] + created: Optional[int] = Field(default_factory=lambda: int(time.time())) + + +@app.get("/v1/models", response_model=ModelList) +async def list_models(): + global model_args + model_card = ModelCard(id="gpt-3.5-turbo") + return ModelList(data=[model_card]) + + +@app.post("/v1/chat/completions", response_model=ChatCompletionResponse) +async def create_chat_completion(request: ChatCompletionRequest): + global model, tokenizer + + if request.messages[-1].role != "user": + raise HTTPException(status_code=400, detail="Invalid request") + query = request.messages[-1].content + + prev_messages = request.messages[:-1] + if len(prev_messages) > 0 and prev_messages[0].role == "system": + query = prev_messages.pop(0).content + query + + history = [] + if len(prev_messages) % 2 == 0: + for i in range(0, len(prev_messages), 2): + if prev_messages[i].role == "user" and prev_messages[i+1].role == "assistant": + history.append([prev_messages[i].content, prev_messages[i+1].content]) + + if request.stream: + generate = predict(query, history, request.model) + return EventSourceResponse(generate, media_type="text/event-stream") + + response, _ = model.chat(tokenizer, query, history=history) + choice_data = ChatCompletionResponseChoice( + index=0, + message=ChatMessage(role="assistant", content=response), + finish_reason="stop" + ) + + return ChatCompletionResponse(model=request.model, choices=[choice_data], object="chat.completion") + + +async def predict(query: str, history: List[List[str]], model_id: str): + global model, tokenizer + + choice_data = ChatCompletionResponseStreamChoice( + index=0, + delta=DeltaMessage(role="assistant"), + finish_reason=None + ) + chunk = ChatCompletionResponse(model=model_id, choices=[choice_data], object="chat.completion.chunk") + yield "{}".format(chunk.json(exclude_unset=True, ensure_ascii=False)) + + current_length = 0 + + for new_response, _ in model.stream_chat(tokenizer, query, history): + if len(new_response) == current_length: + continue + + new_text = new_response[current_length:] + current_length = len(new_response) + + choice_data = ChatCompletionResponseStreamChoice( + index=0, + delta=DeltaMessage(content=new_text), + finish_reason=None + ) + chunk = ChatCompletionResponse(model=model_id, choices=[choice_data], object="chat.completion.chunk") + yield "{}".format(chunk.json(exclude_unset=True, ensure_ascii=False)) + + + choice_data = ChatCompletionResponseStreamChoice( + index=0, + delta=DeltaMessage(), + finish_reason="stop" + ) + chunk = ChatCompletionResponse(model=model_id, choices=[choice_data], object="chat.completion.chunk") + yield "{}".format(chunk.json(exclude_unset=True, ensure_ascii=False)) + yield '[DONE]' + + + +if __name__ == "__main__": + tokenizer = AutoTokenizer.from_pretrained("THUDM/chatglm2-6b", trust_remote_code=True) + model = AutoModel.from_pretrained("THUDM/chatglm2-6b", trust_remote_code=True).cuda() + # 多显卡支持，使用下面两行代替上面一行，将num_gpus改为你实际的显卡数量 + # from utils import load_model_on_gpus + # model = load_model_on_gpus("THUDM/chatglm2-6b", num_gpus=2) + model.eval() + + uvicorn.run(app, host='0.0.0.0', port=8000, workers=1) diff --git a/PyTorch/built-in/foundation/CodeGeeX2/ptuning/arguments.py b/PyTorch/built-in/foundation/CodeGeeX2/ptuning/arguments.py new file mode 100644 index 0000000000000000000000000000000000000000..c6fdcef07f9840cbfc6399f81e9d3fcc5c2900a8 --- /dev/null +++ b/PyTorch/built-in/foundation/CodeGeeX2/ptuning/arguments.py @@ -0,0 +1,224 @@ +from dataclasses import dataclass, field +from typing import Optional + + +@dataclass +class ModelArguments: + """ + Arguments pertaining to which model/config/tokenizer we are going to fine-tune from. + """ + + model_name_or_path: str = field( + metadata={"help": "Path to pretrained model or model identifier from huggingface.co/models"} + ) + ptuning_checkpoint: str = field( + default=None, metadata={"help": "Path to p-tuning v2 checkpoints"} + ) + config_name: Optional[str] = field( + default=None, metadata={"help": "Pretrained config name or path if not the same as model_name"} + ) + tokenizer_name: Optional[str] = field( + default=None, metadata={"help": "Pretrained tokenizer name or path if not the same as model_name"} + ) + cache_dir: Optional[str] = field( + default=None, + metadata={"help": "Where to store the pretrained models downloaded from huggingface.co"}, + ) + use_fast_tokenizer: bool = field( + default=True, + metadata={"help": "Whether to use one of the fast tokenizer (backed by the tokenizers library) or not."}, + ) + model_revision: str = field( + default="main", + metadata={"help": "The specific model version to use (can be a branch name, tag name or commit id)."}, + ) + use_auth_token: bool = field( + default=False, + metadata={ + "help": ( + "Will use the token generated when running `huggingface-cli login` (necessary to use this script " + "with private models)." + ) + }, + ) + resize_position_embeddings: Optional[bool] = field( + default=None, + metadata={ + "help": ( + "Whether to automatically resize the position embeddings if `max_source_length` exceeds " + "the model's position embeddings." + ) + }, + ) + quantization_bit: Optional[int] = field( + default=None + ) + pre_seq_len: Optional[int] = field( + default=None + ) + prefix_projection: bool = field( + default=False + ) + + +@dataclass +class DataTrainingArguments: + """ + Arguments pertaining to what data we are going to input our model for training and eval. + """ + + lang: Optional[str] = field(default=None, metadata={"help": "Language id for summarization."}) + + dataset_name: Optional[str] = field( + default=None, metadata={"help": "The name of the dataset to use (via the datasets library)."} + ) + dataset_config_name: Optional[str] = field( + default=None, metadata={"help": "The configuration name of the dataset to use (via the datasets library)."} + ) + prompt_column: Optional[str] = field( + default=None, + metadata={"help": "The name of the column in the datasets containing the full texts (for summarization)."}, + ) + response_column: Optional[str] = field( + default=None, + metadata={"help": "The name of the column in the datasets containing the summaries (for summarization)."}, + ) + history_column: Optional[str] = field( + default=None, + metadata={"help": "The name of the column in the datasets containing the history of chat."}, + ) + train_file: Optional[str] = field( + default=None, metadata={"help": "The input training data file (a jsonlines or csv file)."} + ) + validation_file: Optional[str] = field( + default=None, + metadata={ + "help": ( + "An optional input evaluation data file to evaluate the metrics (rouge) on (a jsonlines or csv file)." + ) + }, + ) + test_file: Optional[str] = field( + default=None, + metadata={ + "help": "An optional input test data file to evaluate the metrics (rouge) on (a jsonlines or csv file)." + }, + ) + overwrite_cache: bool = field( + default=False, metadata={"help": "Overwrite the cached training and evaluation sets"} + ) + preprocessing_num_workers: Optional[int] = field( + default=None, + metadata={"help": "The number of processes to use for the preprocessing."}, + ) + max_source_length: Optional[int] = field( + default=1024, + metadata={ + "help": ( + "The maximum total input sequence length after tokenization. Sequences longer " + "than this will be truncated, sequences shorter will be padded." + ) + }, + ) + max_target_length: Optional[int] = field( + default=128, + metadata={ + "help": ( + "The maximum total sequence length for target text after tokenization. Sequences longer " + "than this will be truncated, sequences shorter will be padded." + ) + }, + ) + val_max_target_length: Optional[int] = field( + default=None, + metadata={ + "help": ( + "The maximum total sequence length for validation target text after tokenization. Sequences longer " + "than this will be truncated, sequences shorter will be padded. Will default to `max_target_length`." + "This argument is also used to override the ``max_length`` param of ``model.generate``, which is used " + "during ``evaluate`` and ``predict``." + ) + }, + ) + pad_to_max_length: bool = field( + default=False, + metadata={ + "help": ( + "Whether to pad all samples to model maximum sentence length. " + "If False, will pad the samples dynamically when batching to the maximum length in the batch. More " + "efficient on GPU but very bad for TPU." + ) + }, + ) + max_train_samples: Optional[int] = field( + default=None, + metadata={ + "help": ( + "For debugging purposes or quicker training, truncate the number of training examples to this " + "value if set." + ) + }, + ) + max_eval_samples: Optional[int] = field( + default=None, + metadata={ + "help": ( + "For debugging purposes or quicker training, truncate the number of evaluation examples to this " + "value if set." + ) + }, + ) + max_predict_samples: Optional[int] = field( + default=None, + metadata={ + "help": ( + "For debugging purposes or quicker training, truncate the number of prediction examples to this " + "value if set." + ) + }, + ) + num_beams: Optional[int] = field( + default=None, + metadata={ + "help": ( + "Number of beams to use for evaluation. This argument will be passed to ``model.generate``, " + "which is used during ``evaluate`` and ``predict``." + ) + }, + ) + ignore_pad_token_for_loss: bool = field( + default=True, + metadata={ + "help": "Whether to ignore the tokens corresponding to padded labels in the loss computation or not." + }, + ) + source_prefix: Optional[str] = field( + default="", metadata={"help": "A prefix to add before every source text (useful for T5 models)."} + ) + + forced_bos_token: Optional[str] = field( + default=None, + metadata={ + "help": ( + "The token to force as the first generated token after the decoder_start_token_id." + "Useful for multilingual models like mBART where the first generated token" + "needs to be the target language token (Usually it is the target language token)" + ) + }, + ) + + + + def __post_init__(self): + if self.dataset_name is None and self.train_file is None and self.validation_file is None and self.test_file is None: + raise ValueError("Need either a dataset name or a training/validation/test file.") + else: + if self.train_file is not None: + extension = self.train_file.split(".")[-1] + assert extension in ["csv", "json"], "`train_file` should be a csv or a json file." + if self.validation_file is not None: + extension = self.validation_file.split(".")[-1] + assert extension in ["csv", "json"], "`validation_file` should be a csv or a json file." + if self.val_max_target_length is None: + self.val_max_target_length = self.max_target_length + diff --git a/PyTorch/built-in/foundation/CodeGeeX2/ptuning/code_alpaca_20k.json b/PyTorch/built-in/foundation/CodeGeeX2/ptuning/code_alpaca_20k.json new file mode 100644 index 0000000000000000000000000000000000000000..8c4a7facc801b2293db46898fc166a737c41a9e7 --- /dev/null +++ b/PyTorch/built-in/foundation/CodeGeeX2/ptuning/code_alpaca_20k.json @@ -0,0 +1 @@ +{"payload":{"allShortcutsEnabled":false,"fileTree":{"data":{"items":[{"name":"code_alpaca_20k.json","path":"data/code_alpaca_20k.json","contentType":"file"},{"name":"code_alpaca_2k.json","path":"data/code_alpaca_2k.json","contentType":"file"},{"name":"new_codealpaca.json","path":"data/new_codealpaca.json","contentType":"file"},{"name":"rosetta_alpaca.json","path":"data/rosetta_alpaca.json","contentType":"file"}],"totalCount":4},"":{"items":[{"name":"data","path":"data","contentType":"directory"},{"name":".gitignore","path":".gitignore","contentType":"file"},{"name":"DATA_LICENSE","path":"DATA_LICENSE","contentType":"file"},{"name":"LICENSE","path":"LICENSE","contentType":"file"},{"name":"README.md","path":"README.md","contentType":"file"},{"name":"convert_to_hf.py","path":"convert_to_hf.py","contentType":"file"},{"name":"ds_config.json","path":"ds_config.json","contentType":"file"},{"name":"generate_instruction.py","path":"generate_instruction.py","contentType":"file"},{"name":"nolora.py","path":"nolora.py","contentType":"file"},{"name":"prompt.txt","path":"prompt.txt","contentType":"file"},{"name":"requirements.txt","path":"requirements.txt","contentType":"file"},{"name":"seed_tasks.jsonl","path":"seed_tasks.jsonl","contentType":"file"},{"name":"train.py","path":"train.py","contentType":"file"},{"name":"utils.py","path":"utils.py","contentType":"file"}],"totalCount":14}},"fileTreeProcessingTime":7.647501,"foldersToFetch":[],"reducedMotionEnabled":null,"repo":{"id":617429416,"defaultBranch":"master","name":"codealpaca","ownerLogin":"sahil280114","currentUserCanPush":false,"isFork":false,"isEmpty":false,"createdAt":"2023-03-22T11:28:54.000Z","ownerAvatar":"https://avatars.githubusercontent.com/u/22678055?v=4","public":true,"private":false,"isOrgOwned":false},"symbolsExpanded":false,"treeExpanded":true,"refInfo":{"name":"master","listCacheKey":"v0:1679484559.0","canEdit":false,"refType":"branch","currentOid":"2f78ddc5c682ed6738ad092bbbfa59ba915afcb0"},"path":"data/code_alpaca_20k.json","currentUser":null,"blob":{"rawLines":[],"stylingDirectives":[],"csv":null,"csvError":null,"dependabotInfo":{"showConfigurationBanner":false,"configFilePath":null,"networkDependabotPath":"/sahil280114/codealpaca/network/updates","dismissConfigurationNoticePath":"/settings/dismiss-notice/dependabot_configuration_notice","configurationNoticeDismissed":null,"repoAlertsPath":"/sahil280114/codealpaca/security/dependabot","repoSecurityAndAnalysisPath":"/sahil280114/codealpaca/settings/security_analysis","repoOwnerIsOrg":false,"currentUserCanAdminRepo":false},"displayName":"code_alpaca_20k.json","displayUrl":"https://github.com/sahil280114/codealpaca/blob/master/data/code_alpaca_20k.json?raw=true","headerInfo":{"blobSize":"7.68 MB","deleteInfo":{"deleteTooltip":"You must be signed in to make or propose changes"},"editInfo":{"editTooltip":"You must be signed in to make or propose changes"},"ghDesktopPath":"https://desktop.github.com","gitLfsPath":null,"onBranch":true,"shortPath":"4599591","siteNavLoginPath":"/login?return_to=https%3A%2F%2Fgithub.com%2Fsahil280114%2Fcodealpaca%2Fblob%2Fmaster%2Fdata%2Fcode_alpaca_20k.json","isCSV":false,"isRichtext":false,"toc":null,"lineInfo":{"truncatedLoc":null,"truncatedSloc":null},"mode":"file"},"image":false,"isCodeownersFile":null,"isPlain":false,"isValidLegacyIssueTemplate":false,"issueTemplateHelpUrl":"https://docs.github.com/articles/about-issue-and-pull-request-templates","issueTemplate":null,"discussionTemplate":null,"language":null,"languageID":null,"large":true,"loggedIn":false,"newDiscussionPath":"/sahil280114/codealpaca/discussions/new","newIssuePath":"/sahil280114/codealpaca/issues/new","planSupportInfo":{"repoIsFork":null,"repoOwnedByCurrentUser":null,"requestFullPath":"/sahil280114/codealpaca/blob/master/data/code_alpaca_20k.json","showFreeOrgGatedFeatureMessage":null,"showPlanSupportBanner":null,"upgradeDataAttributes":null,"upgradePath":null},"publishBannersInfo":{"dismissActionNoticePath":"/settings/dismiss-notice/publish_action_from_dockerfile","dismissStackNoticePath":"/settings/dismiss-notice/publish_stack_from_file","releasePath":"/sahil280114/codealpaca/releases/new?marketplace=true","showPublishActionBanner":false,"showPublishStackBanner":false},"renderImageOrRaw":true,"richText":null,"renderedFileInfo":null,"shortPath":null,"tabSize":8,"topBannersInfo":{"overridingGlobalFundingFile":false,"globalPreferredFundingPath":null,"repoOwner":"sahil280114","repoName":"codealpaca","showInvalidCitationWarning":false,"citationHelpUrl":"https://docs.github.com/en/github/creating-cloning-and-archiving-repositories/creating-a-repository-on-github/about-citation-files","showDependabotConfigurationBanner":false,"actionsOnboardingTip":null},"truncated":true,"viewable":false,"workflowRedirectUrl":null,"symbols":{"timedOut":true,"notAnalyzed":true,"symbols":[],"error":{"code":"invalid_argument","msg":"content required","meta":{}}}},"copilotInfo":null,"csrf_tokens":{"/sahil280114/codealpaca/branches":{"post":"kFWsMRsRy3fnv-Qvn5wiR-hfJ-mbHRYQgLAiwh98Z7SQxOEPbIYyzdyal6KZ7F1lM2IbrzKtDTjb6b0DPojKZg"},"/repos/preferences":{"post":"vfcrddoss6RoohElb9S4bs2UkZneuZ4WdZL-aqeKK74uBXNAMNQwEPoWCoh5EThbLISBikMsVU1qtQsPKq3mEg"}}},"title":"codealpaca/data/code_alpaca_20k.json at master · sahil280114/codealpaca"} \ No newline at end of file diff --git a/PyTorch/built-in/foundation/CodeGeeX2/ptuning/deepspeed.json b/PyTorch/built-in/foundation/CodeGeeX2/ptuning/deepspeed.json new file mode 100644 index 0000000000000000000000000000000000000000..95998c1705f1126fd5fa21b75e1d0994b3dc052b --- /dev/null +++ b/PyTorch/built-in/foundation/CodeGeeX2/ptuning/deepspeed.json @@ -0,0 +1,21 @@ +{ + "train_micro_batch_size_per_gpu": "auto", + "zero_allow_untested_optimizer": true, + "fp16": { + "enabled": "auto", + "loss_scale": 0, + "initial_scale_power": 16, + "loss_scale_window": 1000, + "hysteresis": 2, + "min_loss_scale": 1 + }, + "zero_optimization": { + "stage": 2, + "allgather_partitions": true, + "allgather_bucket_size": 5e8, + "overlap_comm": true, + "reduce_scatter": true, + "reduce_bucket_size": 5e8, + "contiguous_gradients" : true + } +} \ No newline at end of file diff --git a/PyTorch/built-in/foundation/CodeGeeX2/ptuning/ds_train_finetune.sh b/PyTorch/built-in/foundation/CodeGeeX2/ptuning/ds_train_finetune.sh new file mode 100644 index 0000000000000000000000000000000000000000..e57bd16ebfeaa0b525434e8faf5e2af9f03c8b3b --- /dev/null +++ b/PyTorch/built-in/foundation/CodeGeeX2/ptuning/ds_train_finetune.sh @@ -0,0 +1,29 @@ + +#LR=1e-4 +LR=2e-5 + +MASTER_PORT=$(shuf -n 1 -i 10000-65535) +source env_npu.sh +#deepspeed --num_gpus=8 --master_port $MASTER_PORT main.py \ +deepspeed --num_gpus=8 --master_port $MASTER_PORT main_without_tokenizer.py \ + --deepspeed deepspeed.json \ + --do_train \ + --train_file train.json \ + --test_file dev.json \ + --prompt_column content \ + --response_column summary \ + --model_name_or_path /home/h00638954/codegeex2-6b/ \ + --output_dir ./output/adgen-chatglm2-6b-ft-$LR \ + --overwrite_output_dir \ + --max_source_length 2048 \ + --max_target_length 2048 \ + --per_device_train_batch_size 4 \ + --per_device_eval_batch_size 1 \ + --gradient_accumulation_steps 8 \ + --predict_with_generate \ + --max_steps 100 \ + --logging_steps 1 \ + --save_steps 100 \ + --learning_rate $LR \ + --fp16 + diff --git a/PyTorch/built-in/foundation/CodeGeeX2/ptuning/env_npu.sh b/PyTorch/built-in/foundation/CodeGeeX2/ptuning/env_npu.sh new file mode 100644 index 0000000000000000000000000000000000000000..dd812557560e18c0c6452fa3a015b7381fb7ae59 --- /dev/null +++ b/PyTorch/built-in/foundation/CodeGeeX2/ptuning/env_npu.sh @@ -0,0 +1,70 @@ +#!/bin/bash +CANN_INSTALL_PATH_CONF='/etc/Ascend/ascend_cann_install.info' +#source /home/h00638954/cann_7/ascend-toolkit/set_env.sh +source /home/yrl/yrl_cann/ascend-toolkit/set_env.sh +#source /home/data/yaohan/cann0729/ascend/ascend-toolkit/set_env.sh +#source "/home/data/yaohan/0725/ascend/ascend-toolkit/set_env.sh" +#source '/home/wzw/RC2_0725/ascend-toolkit/set_env.sh' +#source /home/data/yaohan/cann_0705/cann/ascend-toolkit/set_env.sh +#source "/home/data/yaohan/cann0624/cann/ascend-toolkit/set_env.sh" +#设置device侧日志登记为error +msnpureport -g error -d 0 +msnpureport -g error -d 1 +msnpureport -g error -d 2 +msnpureport -g error -d 3 +msnpureport -g error -d 4 +msnpureport -g error -d 5 +msnpureport -g error -d 6 +msnpureport -g error -d 7 + +#关闭Device侧Event日志 +msnpureport -e disable +#判断是否溢出 +export INF_NAN_MODE_ENABLE=0 +#将Host日志输出到串口,0-关闭/1-开启 +export ASCEND_SLOG_PRINT_TO_STDOUT=0 +#设置默认日志级别,0-debug/1-info/2-warning/3-error +export ASCEND_GLOBAL_LOG_LEVEL=3 +#设置Event日志开启标志,0-关闭/1-开启 +export ASCEND_GLOBAL_EVENT_ENABLE=0 +#设置是否开启taskque,0-关闭/1-开启 +export TASK_QUEUE_ENABLE=1 +#设置是否开启PTCopy,0-关闭/1-开启 +export PTCOPY_ENABLE=1 +#设置是否开启2个非连续combined标志,0-关闭/1-开启 +export COMBINED_ENABLE=1 +#设置特殊场景是否需要重新编译,不需要修改 +export DYNAMIC_OP="ADD#MUL" +#HCCL白名单开关,1-关闭/0-开启 +export HCCL_WHITELIST_DISABLE=1 +#设置HCCL超时时间 +export HCCL_CONNECT_TIMEOUT=7200 +export ASCEND_WAIT_NO_PREP_PATH=0 +#export HCCL_OP_BASE_FFTS_MODE_ENABLE=TRUE + + + +ulimit -SHn 512000 + +path_lib=$(python3.7 -c """ +import sys +import re +result='' +for index in range(len(sys.path)): + match_sit = re.search('-packages', sys.path[index]) + if match_sit is not None: + match_lib = re.search('lib', sys.path[index]) + + if match_lib is not None: + end=match_lib.span()[1] + result += sys.path[index][0:end] + ':' + + result+=sys.path[index] + '/torch/lib:' +print(result)""" +) + +echo ${path_lib} + +export LD_LIBRARY_PATH=/usr/local/python3.7.5/lib/:${path_lib}:$LD_LIBRARY_PATH +export HCCL_WHITELIST_DISABLE=1 +#export HCCL_IF_IP=$(hostname -I |awk '{print $1}') diff --git a/PyTorch/built-in/foundation/CodeGeeX2/ptuning/main.py b/PyTorch/built-in/foundation/CodeGeeX2/ptuning/main.py new file mode 100644 index 0000000000000000000000000000000000000000..1df82b4ec5138b856fcda14a12ff6659621f3e56 --- /dev/null +++ b/PyTorch/built-in/foundation/CodeGeeX2/ptuning/main.py @@ -0,0 +1,419 @@ +#!/usr/bin/env python +# coding=utf-8 +# Copyright 2021 The HuggingFace Team. 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. +""" +Fine-tuning the library models for sequence to sequence. +""" +# You can also adapt this script on your own sequence to sequence task. Pointers for this are left as comments. + +import logging +import os +import sys +import json + +import numpy as np +from datasets import load_dataset +import jieba +from rouge_chinese import Rouge +from nltk.translate.bleu_score import sentence_bleu, SmoothingFunction +import torch +import torch_npu +import deepspeed_npu +from torch_npu.contrib import transfer_to_npu + +torch.npu.set_compile_mode(jit_compile=True) +option = {"NPU_FUZZY_COMPILE_BLACKLIST":"Tril,LayerNormGrad"} +torch.npu.set_option(option) + + +import transformers +from transformers import ( + AutoConfig, + AutoModel, + AutoTokenizer, + DataCollatorForSeq2Seq, + HfArgumentParser, + Seq2SeqTrainingArguments, + set_seed, +) +from trainer_seq2seq import Seq2SeqTrainer + +from arguments import ModelArguments, DataTrainingArguments + +logger = logging.getLogger(__name__) + +def main(): + parser = HfArgumentParser((ModelArguments, DataTrainingArguments, Seq2SeqTrainingArguments)) + if len(sys.argv) == 2 and sys.argv[1].endswith(".json"): + # If we pass only one argument to the script and it's the path to a json file, + # let's parse it to get our arguments. + model_args, data_args, training_args = parser.parse_json_file(json_file=os.path.abspath(sys.argv[1])) + else: + model_args, data_args, training_args = parser.parse_args_into_dataclasses() + + # Setup logging + logging.basicConfig( + format="%(asctime)s - %(levelname)s - %(name)s - %(message)s", + datefmt="%m/%d/%Y %H:%M:%S", + handlers=[logging.StreamHandler(sys.stdout)], + ) + + if training_args.should_log: + # The default of training_args.log_level is passive, so we set log level at info here to have that default. + transformers.utils.logging.set_verbosity_info() + + log_level = training_args.get_process_log_level() + logger.setLevel(log_level) + # datasets.utils.logging.set_verbosity(log_level) + transformers.utils.logging.set_verbosity(log_level) + transformers.utils.logging.enable_default_handler() + transformers.utils.logging.enable_explicit_format() + + # Log on each process the small summary: + logger.warning( + f"Process rank: {training_args.local_rank}, device: {training_args.device}, n_gpu: {training_args.n_gpu}" + + f"distributed training: {bool(training_args.local_rank != -1)}, 16-bits training: {training_args.fp16}" + ) + logger.info(f"Training/evaluation parameters {training_args}") + + # Set seed before initializing model. + set_seed(training_args.seed) + + # Load dataset + data_files = {} + if data_args.train_file is not None: + data_files["train"] = data_args.train_file + extension = data_args.train_file.split(".")[-1] + if data_args.validation_file is not None: + data_files["validation"] = data_args.validation_file + extension = data_args.validation_file.split(".")[-1] + if data_args.test_file is not None: + data_files["test"] = data_args.test_file + extension = data_args.test_file.split(".")[-1] + + raw_datasets = load_dataset( + extension, + data_files=data_files, + cache_dir=model_args.cache_dir, + use_auth_token=True if model_args.use_auth_token else None, + ) + + # Load pretrained model and tokenizer + config = AutoConfig.from_pretrained(model_args.model_name_or_path, trust_remote_code=True) + config.pre_seq_len = model_args.pre_seq_len + config.prefix_projection = model_args.prefix_projection + + tokenizer = AutoTokenizer.from_pretrained(model_args.model_name_or_path, trust_remote_code=True) + + if model_args.ptuning_checkpoint is not None: + # Evaluation + # Loading extra state dict of prefix encoder + model = AutoModel.from_pretrained(model_args.model_name_or_path, config=config, trust_remote_code=True) + prefix_state_dict = torch.load(os.path.join(model_args.ptuning_checkpoint, "pytorch_model.bin")) + new_prefix_state_dict = {} + for k, v in prefix_state_dict.items(): + if k.startswith("transformer.prefix_encoder."): + new_prefix_state_dict[k[len("transformer.prefix_encoder."):]] = v + model.transformer.prefix_encoder.load_state_dict(new_prefix_state_dict) + else: + model = AutoModel.from_pretrained(model_args.model_name_or_path, config=config, trust_remote_code=True) + + if model_args.quantization_bit is not None: + print(f"Quantized to {model_args.quantization_bit} bit") + model = model.quantize(model_args.quantization_bit) + if model_args.pre_seq_len is not None: + # P-tuning v2 + model = model.half() + model.transformer.prefix_encoder.float() + else: + # Finetune + model = model.float() + + prefix = data_args.source_prefix if data_args.source_prefix is not None else "" + + # Preprocessing the datasets. + # We need to tokenize inputs and targets. + if training_args.do_train: + column_names = raw_datasets["train"].column_names + elif training_args.do_eval: + column_names = raw_datasets["validation"].column_names + elif training_args.do_predict: + column_names = raw_datasets["test"].column_names + else: + logger.info("There is nothing to do. Please pass `do_train`, `do_eval` and/or `do_predict`.") + return + + # Get the column names for input/target. + prompt_column = data_args.prompt_column + response_column = data_args.response_column + history_column = data_args.history_column + + # Temporarily set max_target_length for training. + max_target_length = data_args.max_target_length + + def preprocess_function_eval(examples): + inputs, targets = [], [] + for i in range(len(examples[prompt_column])): + if examples[prompt_column][i] and examples[response_column][i]: + query = examples[prompt_column][i] + history = examples[history_column][i] if history_column is not None else None + prompt = tokenizer.build_prompt(query, history) + inputs.append(prompt) + targets.append(examples[response_column][i]) + + inputs = [prefix + inp for inp in inputs] + model_inputs = tokenizer(inputs, max_length=data_args.max_source_length, truncation=True, padding=True) + labels = tokenizer(text_target=targets, max_length=max_target_length, truncation=True) + + if data_args.ignore_pad_token_for_loss: + labels["input_ids"] = [ + [(l if l != tokenizer.pad_token_id else -100) for l in label] for label in labels["input_ids"] + ] + model_inputs["labels"] = labels["input_ids"] + + return model_inputs + + def preprocess_function_train(examples): + max_seq_length = data_args.max_source_length + data_args.max_target_length + 1 + + model_inputs = { + "input_ids": [], + "labels": [], + } + for i in range(len(examples[prompt_column])): + if examples[prompt_column][i] and examples[response_column][i]: + query, answer = examples[prompt_column][i], examples[response_column][i] + + history = examples[history_column][i] if history_column is not None else None + prompt = tokenizer.build_prompt(query, history) + + prompt = prefix + prompt + a_ids = tokenizer.encode(text=prompt, add_special_tokens=True, truncation=True, + max_length=data_args.max_source_length) + b_ids = tokenizer.encode(text=answer, add_special_tokens=False, truncation=True, + max_length=data_args.max_target_length) + + context_length = len(a_ids) + input_ids = a_ids + b_ids + [tokenizer.eos_token_id] + labels = [tokenizer.pad_token_id] * context_length + b_ids + [tokenizer.eos_token_id] + + pad_len = max_seq_length - len(input_ids) + input_ids = input_ids + [tokenizer.pad_token_id] * pad_len + labels = labels + [tokenizer.pad_token_id] * pad_len + if data_args.ignore_pad_token_for_loss: + labels = [(l if l != tokenizer.pad_token_id else -100) for l in labels] + + model_inputs["input_ids"].append(input_ids) + model_inputs["labels"].append(labels) + + return model_inputs + + def print_dataset_example(example): + print("input_ids", example["input_ids"]) + print("inputs", tokenizer.decode(example["input_ids"])) + print("label_ids", example["labels"]) + print("labels", tokenizer.decode(example["labels"])) + + if training_args.do_train: + if "train" not in raw_datasets: + raise ValueError("--do_train requires a train dataset") + train_dataset = raw_datasets["train"] + if data_args.max_train_samples is not None: + max_train_samples = min(len(train_dataset), data_args.max_train_samples) + train_dataset = train_dataset.select(range(max_train_samples)) + with training_args.main_process_first(desc="train dataset map pre-processing"): + train_dataset = train_dataset.map( + preprocess_function_train, + batched=True, + num_proc=data_args.preprocessing_num_workers, + remove_columns=column_names, + load_from_cache_file=not data_args.overwrite_cache, + desc="Running tokenizer on train dataset", + ) + print_dataset_example(train_dataset[0]) + + if training_args.do_eval: + max_target_length = data_args.val_max_target_length + if "validation" not in raw_datasets: + raise ValueError("--do_eval requires a validation dataset") + eval_dataset = raw_datasets["validation"] + if data_args.max_eval_samples is not None: + max_eval_samples = min(len(eval_dataset), data_args.max_eval_samples) + eval_dataset = eval_dataset.select(range(max_eval_samples)) + with training_args.main_process_first(desc="validation dataset map pre-processing"): + eval_dataset = eval_dataset.map( + preprocess_function_eval, + batched=True, + num_proc=data_args.preprocessing_num_workers, + remove_columns=column_names, + load_from_cache_file=not data_args.overwrite_cache, + desc="Running tokenizer on validation dataset", + ) + print_dataset_example(eval_dataset[0]) + + if training_args.do_predict: + max_target_length = data_args.val_max_target_length + if "test" not in raw_datasets: + raise ValueError("--do_predict requires a test dataset") + predict_dataset = raw_datasets["test"] + if data_args.max_predict_samples is not None: + max_predict_samples = min(len(predict_dataset), data_args.max_predict_samples) + predict_dataset = predict_dataset.select(range(max_predict_samples)) + with training_args.main_process_first(desc="prediction dataset map pre-processing"): + predict_dataset = predict_dataset.map( + preprocess_function_eval, + batched=True, + num_proc=data_args.preprocessing_num_workers, + remove_columns=column_names, + load_from_cache_file=not data_args.overwrite_cache, + desc="Running tokenizer on prediction dataset", + ) + print_dataset_example(predict_dataset[0]) + + # Data collator + label_pad_token_id = -100 if data_args.ignore_pad_token_for_loss else tokenizer.pad_token_id + data_collator = DataCollatorForSeq2Seq( + tokenizer, + model=model, + label_pad_token_id=label_pad_token_id, + pad_to_multiple_of=None, + padding=False + ) + + # Metric + def compute_metrics(eval_preds): + preds, labels = eval_preds + if isinstance(preds, tuple): + preds = preds[0] + decoded_preds = tokenizer.batch_decode(preds, skip_special_tokens=True) + if data_args.ignore_pad_token_for_loss: + # Replace -100 in the labels as we can't decode them. + labels = np.where(labels != -100, labels, tokenizer.pad_token_id) + decoded_labels = tokenizer.batch_decode(labels, skip_special_tokens=True) + + score_dict = { + "rouge-1": [], + "rouge-2": [], + "rouge-l": [], + "bleu-4": [] + } + for pred, label in zip(decoded_preds, decoded_labels): + hypothesis = list(jieba.cut(pred)) + reference = list(jieba.cut(label)) + rouge = Rouge() + scores = rouge.get_scores(' '.join(hypothesis) , ' '.join(reference)) + result = scores[0] + + for k, v in result.items(): + score_dict[k].append(round(v["f"] * 100, 4)) + bleu_score = sentence_bleu([list(label)], list(pred), smoothing_function=SmoothingFunction().method3) + score_dict["bleu-4"].append(round(bleu_score * 100, 4)) + + for k, v in score_dict.items(): + score_dict[k] = float(np.mean(v)) + return score_dict + + # Override the decoding parameters of Seq2SeqTrainer + training_args.generation_max_length = ( + training_args.generation_max_length + if training_args.generation_max_length is not None + else data_args.val_max_target_length + ) + training_args.generation_num_beams = ( + data_args.num_beams if data_args.num_beams is not None else training_args.generation_num_beams + ) + # Initialize our Trainer + trainer = Seq2SeqTrainer( + model=model, + args=training_args, + train_dataset=train_dataset if training_args.do_train else None, + eval_dataset=eval_dataset if training_args.do_eval else None, + tokenizer=tokenizer, + data_collator=data_collator, + compute_metrics=compute_metrics if training_args.predict_with_generate else None, + save_changed=model_args.pre_seq_len is not None + ) + + # Training + if training_args.do_train: + checkpoint = None + if training_args.resume_from_checkpoint is not None: + checkpoint = training_args.resume_from_checkpoint + # elif last_checkpoint is not None: + # checkpoint = last_checkpoint + model.gradient_checkpointing_enable() + model.enable_input_require_grads() + train_result = trainer.train(resume_from_checkpoint=checkpoint) + # trainer.save_model() # Saves the tokenizer too for easy upload + + metrics = train_result.metrics + max_train_samples = ( + data_args.max_train_samples if data_args.max_train_samples is not None else len(train_dataset) + ) + metrics["train_samples"] = min(max_train_samples, len(train_dataset)) + + trainer.log_metrics("train", metrics) + trainer.save_metrics("train", metrics) + trainer.save_state() + + # Evaluation + results = {} + max_seq_length = data_args.max_source_length + data_args.max_target_length + 1 + if training_args.do_eval: + logger.info("*** Evaluate ***") + metrics = trainer.evaluate(metric_key_prefix="eval", do_sample=True, top_p=0.7, max_length=max_seq_length, temperature=0.95) + max_eval_samples = data_args.max_eval_samples if data_args.max_eval_samples is not None else len(eval_dataset) + metrics["eval_samples"] = min(max_eval_samples, len(eval_dataset)) + + trainer.log_metrics("eval", metrics) + trainer.save_metrics("eval", metrics) + + if training_args.do_predict: + logger.info("*** Predict ***") + predict_results = trainer.predict(predict_dataset, metric_key_prefix="predict", max_length=max_seq_length, do_sample=True, top_p=0.7, temperature=0.95) + metrics = predict_results.metrics + max_predict_samples = ( + data_args.max_predict_samples if data_args.max_predict_samples is not None else len(predict_dataset) + ) + metrics["predict_samples"] = min(max_predict_samples, len(predict_dataset)) + + trainer.log_metrics("predict", metrics) + trainer.save_metrics("predict", metrics) + + if trainer.is_world_process_zero(): + if training_args.predict_with_generate: + predictions = tokenizer.batch_decode( + predict_results.predictions, skip_special_tokens=True, clean_up_tokenization_spaces=True + ) + predictions = [pred.strip() for pred in predictions] + labels = tokenizer.batch_decode( + predict_results.label_ids, skip_special_tokens=True, clean_up_tokenization_spaces=True + ) + labels = [label.strip() for label in labels] + output_prediction_file = os.path.join(training_args.output_dir, "generated_predictions.txt") + with open(output_prediction_file, "w", encoding="utf-8") as writer: + for p, l in zip(predictions, labels): + res = json.dumps({"labels": l, "predict": p}, ensure_ascii=False) + writer.write(f"{res}\n") + return results + + +def _mp_fn(index): + # For xla_spawn (TPUs) + main() + + +if __name__ == "__main__": + main() diff --git a/PyTorch/built-in/foundation/CodeGeeX2/ptuning/main_without_tokenizer.py b/PyTorch/built-in/foundation/CodeGeeX2/ptuning/main_without_tokenizer.py new file mode 100644 index 0000000000000000000000000000000000000000..38390273aae5dc532ea275c0b91380cf6edb212e --- /dev/null +++ b/PyTorch/built-in/foundation/CodeGeeX2/ptuning/main_without_tokenizer.py @@ -0,0 +1,307 @@ +#!/usr/bin/env python +# coding=utf-8 +# Copyright 2023 Huawei Technologies Co., Ltd +# +# 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. + +# coding=utf-8 +# Copyright 2021 The HuggingFace Team. 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. +""" +Fine-tuning the library models for sequence to sequence. +""" +# You can also adapt this script on your own sequence to sequence task. Pointers for this are left as comments. + +import logging +import os +import sys +import json + +import numpy as np +from datasets import load_dataset, load_from_disk +import jieba +from rouge_chinese import Rouge +from nltk.translate.bleu_score import sentence_bleu, SmoothingFunction +import torch +import torch_npu +import deepspeed_npu +from torch_npu.contrib import transfer_to_npu + +import transformers +from transformers import ( + AutoConfig, + AutoModel, + AutoTokenizer, + DataCollatorForSeq2Seq, + HfArgumentParser, + Seq2SeqTrainingArguments, + set_seed, +) +from trainer_seq2seq import Seq2SeqTrainer + +from arguments import ModelArguments, DataTrainingArguments + +logger = logging.getLogger(__name__) + + +def main(): + torch.use_deterministic_algorithms(True) + parser = HfArgumentParser((ModelArguments, DataTrainingArguments, Seq2SeqTrainingArguments)) + if len(sys.argv) == 2 and sys.argv[1].endswith(".json"): + # If we pass only one argument to the script and it's the path to a json file, + # let's parse it to get our arguments. + model_args, data_args, training_args = parser.parse_json_file(json_file=os.path.abspath(sys.argv[1])) + else: + model_args, data_args, training_args = parser.parse_args_into_dataclasses() + if training_args.do_train: + torch.npu.set_compile_mode(jit_compile=True) + option = {"NPU_FUZZY_COMPILE_BLACKLIST": "Tril,LayerNormGrad"} + torch.npu.set_option(option) + + if training_args.do_predict: + training_args.local_rank = -1 + # Setup logging + logging.basicConfig( + format="%(asctime)s - %(levelname)s - %(name)s - %(message)s", + datefmt="%m/%d/%Y %H:%M:%S", + handlers=[logging.StreamHandler(sys.stdout)], + ) + + if training_args.should_log: + # The default of training_args.log_level is passive, so we set log level at info here to have that default. + transformers.utils.logging.set_verbosity_info() + + log_level = training_args.get_process_log_level() + logger.setLevel(log_level) + # datasets.utils.logging.set_verbosity(log_level) + transformers.utils.logging.set_verbosity(log_level) + transformers.utils.logging.enable_default_handler() + transformers.utils.logging.enable_explicit_format() + + # Log on each process the small summary: + logger.warning( + f"Process rank: {training_args.local_rank}, device: {training_args.device}, n_gpu: {training_args.n_gpu}" + + f"distributed training: {bool(training_args.local_rank != -1)}, 16-bits training: {training_args.fp16}" + ) + logger.info(f"Training/evaluation parameters {training_args}") + + # Set seed before initializing model. + set_seed(training_args.seed) + + # Load pretrained model and tokenizer + config = AutoConfig.from_pretrained(model_args.model_name_or_path, trust_remote_code=True) + config.pre_seq_len = model_args.pre_seq_len + config.prefix_projection = model_args.prefix_projection + + tokenizer = AutoTokenizer.from_pretrained(model_args.model_name_or_path, trust_remote_code=True) + + if model_args.ptuning_checkpoint is not None: + # Evaluation + # Loading extra state dict of prefix encoder + model = AutoModel.from_pretrained(model_args.model_name_or_path, config=config, trust_remote_code=True) + prefix_state_dict = torch.load(os.path.join(model_args.ptuning_checkpoint, "pytorch_model.bin")) + new_prefix_state_dict = {} + for k, v in prefix_state_dict.items(): + if k.startswith("transformer.prefix_encoder."): + new_prefix_state_dict[k[len("transformer.prefix_encoder."):]] = v + model.transformer.prefix_encoder.load_state_dict(new_prefix_state_dict) + else: + model = AutoModel.from_pretrained(model_args.model_name_or_path, config=config, trust_remote_code=True) + + if model_args.quantization_bit is not None: + print(f"Quantized to {model_args.quantization_bit} bit") + model = model.quantize(model_args.quantization_bit) + if model_args.pre_seq_len is not None: + # P-tuning v2 + model = model.half() + model.transformer.prefix_encoder.float() + else: + # Finetune + model = model.float() + + prefix = data_args.source_prefix if data_args.source_prefix is not None else "" + + def print_dataset_example(example): + print("input_ids", example["input_ids"]) + print("inputs", tokenizer.decode(example["input_ids"])) + print("label_ids", example["labels"]) + print("labels", tokenizer.decode(example["labels"])) + + if training_args.do_train: + train_dataset = load_from_disk("train_datasets") + + print_dataset_example(train_dataset[0]) + if data_args.max_train_samples is not None: + max_train_samples = min(len(train_dataset), data_args.max_train_samples) + train_dataset = train_dataset.select(range(max_train_samples)) + + if training_args.do_eval: + eval_dataset = load_from_disk("validation_datasets") + if data_args.max_eval_samples is not None: + max_eval_samples = min(len(eval_dataset), data_args.max_eval_samples) + eval_dataset = eval_dataset.select(range(max_eval_samples)) + + if training_args.do_predict: + predict_dataset = load_from_disk("test_datasets") + if data_args.max_predict_samples is not None: + max_predict_samples = min(len(predict_dataset), data_args.max_predict_samples) + predict_dataset = predict_dataset.select(range(max_predict_samples)) + + # Data collator + label_pad_token_id = -100 if data_args.ignore_pad_token_for_loss else tokenizer.pad_token_id + data_collator = DataCollatorForSeq2Seq( + tokenizer, + model=model, + label_pad_token_id=label_pad_token_id, + pad_to_multiple_of=None, + padding=False + ) + + # Metric + def compute_metrics(eval_preds): + preds, labels = eval_preds + if isinstance(preds, tuple): + preds = preds[0] + decoded_preds = tokenizer.batch_decode(preds, skip_special_tokens=True) + if data_args.ignore_pad_token_for_loss: + # Replace -100 in the labels as we can't decode them. + labels = np.where(labels != -100, labels, tokenizer.pad_token_id) + decoded_labels = tokenizer.batch_decode(labels, skip_special_tokens=True) + + score_dict = { + "rouge-1": [], + "rouge-2": [], + "rouge-l": [], + "bleu-4": [] + } + for pred, label in zip(decoded_preds, decoded_labels): + hypothesis = list(jieba.cut(pred)) + reference = list(jieba.cut(label)) + rouge = Rouge() + scores = rouge.get_scores(' '.join(hypothesis), ' '.join(reference)) + result = scores[0] + + for k, v in result.items(): + score_dict[k].append(round(v["f"] * 100, 4)) + bleu_score = sentence_bleu([list(label)], list(pred), smoothing_function=SmoothingFunction().method3) + score_dict["bleu-4"].append(round(bleu_score * 100, 4)) + + for k, v in score_dict.items(): + score_dict[k] = float(np.mean(v)) + return score_dict + + # Override the decoding parameters of Seq2SeqTrainer + training_args.generation_max_length = ( + training_args.generation_max_length + if training_args.generation_max_length is not None + else data_args.val_max_target_length + ) + training_args.generation_num_beams = ( + data_args.num_beams if data_args.num_beams is not None else training_args.generation_num_beams + ) + # Initialize our Trainer + trainer = Seq2SeqTrainer( + model=model, + args=training_args, + train_dataset=train_dataset if training_args.do_train else None, + eval_dataset=eval_dataset if training_args.do_eval else None, + tokenizer=tokenizer, + data_collator=data_collator, + compute_metrics=compute_metrics if training_args.predict_with_generate else None, + save_changed=model_args.pre_seq_len is not None + ) + + # Training + if training_args.do_train: + checkpoint = None + if training_args.resume_from_checkpoint is not None: + checkpoint = training_args.resume_from_checkpoint + # elif last_checkpoint is not None: + # checkpoint = last_checkpoint + model.gradient_checkpointing_enable() + model.enable_input_require_grads() + train_result = trainer.train(resume_from_checkpoint=checkpoint) + # trainer.save_model() # Saves the tokenizer too for easy upload + + metrics = train_result.metrics + max_train_samples = ( + data_args.max_train_samples if data_args.max_train_samples is not None else len(train_dataset) + ) + metrics["train_samples"] = min(max_train_samples, len(train_dataset)) + + trainer.log_metrics("train", metrics) + trainer.save_metrics("train", metrics) + trainer.save_state() + + # Evaluation + results = {} + max_seq_length = data_args.max_source_length + data_args.max_target_length + 1 + if training_args.do_eval: + logger.info("*** Evaluate ***") + metrics = trainer.evaluate(metric_key_prefix="eval", do_sample=True, top_p=0.7, max_length=max_seq_length, + temperature=0.95) + max_eval_samples = data_args.max_eval_samples if data_args.max_eval_samples is not None else len(eval_dataset) + metrics["eval_samples"] = min(max_eval_samples, len(eval_dataset)) + + trainer.log_metrics("eval", metrics) + trainer.save_metrics("eval", metrics) + + if training_args.do_predict: + logger.info("*** Predict ***") + predict_results = trainer.predict(predict_dataset, metric_key_prefix="predict", max_length=max_seq_length, + do_sample=True, top_p=0.7, temperature=0.95) + metrics = predict_results.metrics + max_predict_samples = ( + data_args.max_predict_samples if data_args.max_predict_samples is not None else len(predict_dataset) + ) + metrics["predict_samples"] = min(max_predict_samples, len(predict_dataset)) + + trainer.log_metrics("predict", metrics) + trainer.save_metrics("predict", metrics) + + if trainer.is_world_process_zero(): + if training_args.predict_with_generate: + predictions = tokenizer.batch_decode( + predict_results.predictions, skip_special_tokens=True, clean_up_tokenization_spaces=True + ) + predictions = [pred.strip() for pred in predictions] + labels = tokenizer.batch_decode( + predict_results.label_ids, skip_special_tokens=True, clean_up_tokenization_spaces=True + ) + labels = [label.strip() for label in labels] + output_prediction_file = os.path.join(training_args.output_dir, "generated_predictions.txt") + with open(output_prediction_file, "w", encoding="utf-8") as writer: + for p, l in zip(predictions, labels): + res = json.dumps({"labels": l, "predict": p}, ensure_ascii=False) + writer.write(f"{res}\n") + return results + + +def _mp_fn(index): + # For xla_spawn (TPUs) + main() + + +if __name__ == "__main__": + main() diff --git a/PyTorch/built-in/foundation/CodeGeeX2/ptuning/preprocess.py b/PyTorch/built-in/foundation/CodeGeeX2/ptuning/preprocess.py new file mode 100644 index 0000000000000000000000000000000000000000..41a3ae052e45761b8f095b7d15b87be14e58a282 --- /dev/null +++ b/PyTorch/built-in/foundation/CodeGeeX2/ptuning/preprocess.py @@ -0,0 +1,197 @@ + +import logging +import os +import sys +import json + +import numpy as np +import pandas as pd +from datasets import load_dataset, load_from_disk +import datasets +import jieba +from rouge_chinese import Rouge +from nltk.translate.bleu_score import sentence_bleu, SmoothingFunction +import torch +import torch_npu +import deepspeed_npu +from torch_npu.contrib import transfer_to_npu + +import transformers +from transformers import ( + AutoConfig, + AutoModel, + AutoTokenizer, + DataCollatorForSeq2Seq, + HfArgumentParser, + Seq2SeqTrainingArguments, + set_seed, +) +from trainer_seq2seq import Seq2SeqTrainer + +from arguments import ModelArguments, DataTrainingArguments + +logger = logging.getLogger(__name__) + +def main(): + parser = HfArgumentParser((ModelArguments, DataTrainingArguments, Seq2SeqTrainingArguments)) + if len(sys.argv) == 2 and sys.argv[1].endswith(".json"): + # If we pass only one argument to the script and it's the path to a json file, + # let's parse it to get our arguments. + model_args, data_args, training_args = parser.parse_json_file(json_file=os.path.abspath(sys.argv[1])) + else: + model_args, data_args, training_args = parser.parse_args_into_dataclasses() + + + + # Set seed before initializing model. + set_seed(training_args.seed) + + # Load dataset + data_files = {} + if data_args.train_file is not None: + data_files["train"] = data_args.train_file + extension = data_args.train_file.split(".")[-1] + if data_args.validation_file is not None: + data_files["validation"] = data_args.validation_file + extension = data_args.validation_file.split(".")[-1] + if data_args.test_file is not None: + data_files["test"] = data_args.test_file + extension = data_args.test_file.split(".")[-1] + + raw_datasets = load_dataset( + extension, + data_files=data_files, + cache_dir=model_args.cache_dir, + use_auth_token=True if model_args.use_auth_token else None, + ) + + # Load pretrained model and tokenizer + config = AutoConfig.from_pretrained(model_args.model_name_or_path, trust_remote_code=True) + config.pre_seq_len = model_args.pre_seq_len + config.prefix_projection = model_args.prefix_projection + + tokenizer = AutoTokenizer.from_pretrained(model_args.model_name_or_path, trust_remote_code=True) + + + + prefix = data_args.source_prefix if data_args.source_prefix is not None else "" + + # Preprocessing the datasets. + # We need to tokenize inputs and targets. + if training_args.do_train: + column_names = raw_datasets["train"].column_names + elif training_args.do_eval: + column_names = raw_datasets["validation"].column_names + elif training_args.do_predict: + column_names = raw_datasets["test"].column_names + else: + logger.info("There is nothing to do. Please pass `do_train`, `do_eval` and/or `do_predict`.") + return + + # Get the column names for input/target. + prompt_column = data_args.prompt_column + response_column = data_args.response_column + history_column = data_args.history_column + + # Temporarily set max_target_length for training. + max_target_length = data_args.max_target_length + + def preprocess_function_eval(examples): + inputs, targets = [], [] + for i in range(len(examples[prompt_column])): + if examples[prompt_column][i] and examples[response_column][i]: + query = examples[prompt_column][i] + history = examples[history_column][i] if history_column is not None else None + prompt = tokenizer.build_prompt(query, history) + inputs.append(prompt) + targets.append(examples[response_column][i]) + + inputs = [prefix + inp for inp in inputs] + model_inputs = tokenizer(inputs, max_length=data_args.max_source_length, truncation=True, padding=True) + labels = tokenizer(text_target=targets, max_length=max_target_length, truncation=True) + + if data_args.ignore_pad_token_for_loss: + labels["input_ids"] = [ + [(l if l != tokenizer.pad_token_id else -100) for l in label] for label in labels["input_ids"] + ] + model_inputs["labels"] = labels["input_ids"] + + return model_inputs + + def preprocess_function_train(examples): + # max_seq_length = data_args.max_source_length + data_args.max_target_length + 1 + max_seq_length = data_args.max_source_length + data_args.max_target_length + + model_inputs = { + "input_ids": [], + "labels": [], + } + for i in range(len(examples[prompt_column])): + if examples[prompt_column][i] and examples[response_column][i]: + query, answer = examples[prompt_column][i], examples[response_column][i] + + history = examples[history_column][i] if history_column is not None else None + prompt = tokenizer.build_prompt(query, history) + + prompt = prefix + prompt + a_ids = tokenizer.encode(text=prompt, add_special_tokens=True, truncation=True, + max_length=data_args.max_source_length) + b_ids = tokenizer.encode(text=answer, add_special_tokens=False, truncation=True, + max_length=data_args.max_target_length) + + context_length = len(a_ids) + # Tag:差不多 a+b+eos; + input_ids = a_ids + b_ids + [tokenizer.eos_token_id] + labels = [tokenizer.pad_token_id] * context_length + b_ids + [tokenizer.eos_token_id] + # print("input_ids.shape, labels.shape ", input_ids.shape, labels.shape) + pad_len = max_seq_length - len(input_ids) + input_ids = input_ids + [tokenizer.pad_token_id] * pad_len + labels = labels + [tokenizer.pad_token_id] * pad_len + # print("2==== input_ids.shape, labels.shape ", input_ids.shape, labels.shape) + if data_args.ignore_pad_token_for_loss: + labels = [(l if l != tokenizer.pad_token_id else -100) for l in labels] + + model_inputs["input_ids"].append(input_ids) + model_inputs["labels"].append(labels) + + return model_inputs + + def process_dataset(phase: str): + if phase not in raw_datasets: + raise ValueError("--do_train requires a train dataset") + train_dataset = raw_datasets[phase] + with training_args.main_process_first(desc=f"{phase} dataset map pre-processing"): + train_dataset = train_dataset.map( + func_dict.get(phase), + batched=True, + num_proc=data_args.preprocessing_num_workers, + remove_columns=column_names, + load_from_cache_file=not data_args.overwrite_cache, + desc=f"Running tokenizer on {phase} dataset", + ) + print("train_dataset.shape", type(train_dataset[0]), len(train_dataset[0])) + model_inputs_datasets = datasets.Dataset.from_pandas(pd.DataFrame(train_dataset)) + model_inputs_datasets.save_to_disk(f"{phase}_datasets") + + + func_dict = {"train": preprocess_function_train, "validation": preprocess_function_eval, + "test": preprocess_function_eval} + if training_args.do_train: + process_dataset("train") + # print_dataset_example(train_dataset[0]) + elif training_args.do_eval: + process_dataset("validation") + elif training_args.do_predict: + process_dataset("test") + + +# def _mp_fn(index): +# # For xla_spawn (TPUs) +# main() + + +if __name__ == "__main__": + main() + # Tag:载入数据 + # data = load_from_disk("model_inputs_datasets") + # print("data", data) diff --git a/PyTorch/built-in/foundation/CodeGeeX2/ptuning/preprocess.sh b/PyTorch/built-in/foundation/CodeGeeX2/ptuning/preprocess.sh new file mode 100644 index 0000000000000000000000000000000000000000..95f77e1b06ffaec31ce7b9e458f07f0f4534d020 --- /dev/null +++ b/PyTorch/built-in/foundation/CodeGeeX2/ptuning/preprocess.sh @@ -0,0 +1,25 @@ +source env_npu.sh +python preprocess.py \ + --do_train \ + --train_file train.json \ + --test_file dev.json \ + --prompt_column PROMPT \ + --response_column ANSWER \ + --model_name_or_path /home/h00638954/codegeex2-6b/ \ + --overwrite_cache \ + --output_dir ./output/adgen-chatglm-6b-ft-$LR \ + --max_source_length 2048 \ + --max_target_length 2048 +:< Dict[str, float]: + """ + Run evaluation and returns metrics. + + The calling script will be responsible for providing a method to compute metrics, as they are task-dependent + (pass it to the init `compute_metrics` argument). + + You can also subclass and override this method to inject custom behavior. + + Args: + eval_dataset (`Dataset`, *optional*): + Pass a dataset if you wish to override `self.eval_dataset`. If it is an [`~datasets.Dataset`], columns + not accepted by the `model.forward()` method are automatically removed. It must implement the `__len__` + method. + ignore_keys (`List[str]`, *optional*): + A list of keys in the output of your model (if it is a dictionary) that should be ignored when + gathering predictions. + metric_key_prefix (`str`, *optional*, defaults to `"eval"`): + An optional prefix to be used as the metrics key prefix. For example the metrics "bleu" will be named + "eval_bleu" if the prefix is `"eval"` (default) + max_length (`int`, *optional*): + The maximum target length to use when predicting with the generate method. + num_beams (`int`, *optional*): + Number of beams for beam search that will be used when predicting with the generate method. 1 means no + beam search. + gen_kwargs: + Additional `generate` specific kwargs. + + Returns: + A dictionary containing the evaluation loss and the potential metrics computed from the predictions. The + dictionary also contains the epoch number which comes from the training state. + """ + + gen_kwargs = gen_kwargs.copy() + if gen_kwargs.get("max_length") is None and gen_kwargs.get("max_new_tokens") is None: + gen_kwargs["max_length"] = self.args.generation_max_length + gen_kwargs["num_beams"] = ( + gen_kwargs["num_beams"] if gen_kwargs.get("num_beams") is not None else self.args.generation_num_beams + ) + self._gen_kwargs = gen_kwargs + + return super().evaluate(eval_dataset, ignore_keys=ignore_keys, metric_key_prefix=metric_key_prefix) + + def predict( + self, + test_dataset: Dataset, + ignore_keys: Optional[List[str]] = None, + metric_key_prefix: str = "test", + **gen_kwargs + ) -> PredictionOutput: + """ + Run prediction and returns predictions and potential metrics. + + Depending on the dataset and your use case, your test dataset may contain labels. In that case, this method + will also return metrics, like in `evaluate()`. + + Args: + test_dataset (`Dataset`): + Dataset to run the predictions on. If it is a [`~datasets.Dataset`], columns not accepted by the + `model.forward()` method are automatically removed. Has to implement the method `__len__` + ignore_keys (`List[str]`, *optional*): + A list of keys in the output of your model (if it is a dictionary) that should be ignored when + gathering predictions. + metric_key_prefix (`str`, *optional*, defaults to `"eval"`): + An optional prefix to be used as the metrics key prefix. For example the metrics "bleu" will be named + "eval_bleu" if the prefix is `"eval"` (default) + max_length (`int`, *optional*): + The maximum target length to use when predicting with the generate method. + num_beams (`int`, *optional*): + Number of beams for beam search that will be used when predicting with the generate method. 1 means no + beam search. + gen_kwargs: + Additional `generate` specific kwargs. + + + + If your predictions or labels have different sequence lengths (for instance because you're doing dynamic + padding in a token classification task) the predictions will be padded (on the right) to allow for + concatenation into one array. The padding index is -100. + + + + Returns: *NamedTuple* A namedtuple with the following keys: + + - predictions (`np.ndarray`): The predictions on `test_dataset`. + - label_ids (`np.ndarray`, *optional*): The labels (if the dataset contained some). + - metrics (`Dict[str, float]`, *optional*): The potential dictionary of metrics (if the dataset contained + labels). + """ + + gen_kwargs = gen_kwargs.copy() + if gen_kwargs.get("max_length") is None and gen_kwargs.get("max_new_tokens") is None: + gen_kwargs["max_length"] = self.args.generation_max_length + gen_kwargs["num_beams"] = ( + gen_kwargs["num_beams"] if gen_kwargs.get("num_beams") is not None else self.args.generation_num_beams + ) + self._gen_kwargs = gen_kwargs + + + return super().predict(test_dataset, ignore_keys=ignore_keys, metric_key_prefix=metric_key_prefix) + + def prediction_step( + self, + model: nn.Module, + inputs: Dict[str, Union[torch.Tensor, Any]], + prediction_loss_only: bool, + ignore_keys: Optional[List[str]] = None, + ) -> Tuple[Optional[float], Optional[torch.Tensor], Optional[torch.Tensor]]: + """ + Perform an evaluation step on `model` using `inputs`. + + Subclass and override to inject custom behavior. + + Args: + model (`nn.Module`): + The model to evaluate. + inputs (`Dict[str, Union[torch.Tensor, Any]]`): + The inputs and targets of the model. + + The dictionary will be unpacked before being fed to the model. Most models expect the targets under the + argument `labels`. Check your model's documentation for all accepted arguments. + prediction_loss_only (`bool`): + Whether or not to return the loss only. + + Return: + Tuple[Optional[float], Optional[torch.Tensor], Optional[torch.Tensor]]: A tuple with the loss, logits and + labels (each being optional). + """ + + if not self.args.predict_with_generate or prediction_loss_only: + return super().prediction_step( + model, inputs, prediction_loss_only=prediction_loss_only, ignore_keys=ignore_keys + ) + + has_labels = "labels" in inputs + inputs = self._prepare_inputs(inputs) + + # XXX: adapt synced_gpus for fairscale as well + gen_kwargs = self._gen_kwargs.copy() + if gen_kwargs.get("max_length") is None and gen_kwargs.get("max_new_tokens") is None: + gen_kwargs["max_length"] = self.model.config.max_length + gen_kwargs["num_beams"] = ( + gen_kwargs["num_beams"] if gen_kwargs.get("num_beams") is not None else self.model.config.num_beams + ) + default_synced_gpus = True if is_deepspeed_zero3_enabled() else False + gen_kwargs["synced_gpus"] = ( + gen_kwargs["synced_gpus"] if gen_kwargs.get("synced_gpus") is not None else default_synced_gpus + ) + + if "attention_mask" in inputs: + gen_kwargs["attention_mask"] = inputs.get("attention_mask", None) + if "position_ids" in inputs: + gen_kwargs["position_ids"] = inputs.get("position_ids", None) + if "global_attention_mask" in inputs: + gen_kwargs["global_attention_mask"] = inputs.get("global_attention_mask", None) + + # prepare generation inputs + # some encoder-decoder models can have varying encoder's and thus + # varying model input names + if hasattr(self.model, "encoder") and self.model.encoder.main_input_name != self.model.main_input_name: + generation_inputs = inputs[self.model.encoder.main_input_name] + else: + generation_inputs = inputs[self.model.main_input_name] + + gen_kwargs["input_ids"] = generation_inputs + generated_tokens = self.model.generate(**gen_kwargs) + generated_tokens = generated_tokens[:, generation_inputs.size()[-1]:] + + # in case the batch is shorter than max length, the output should be padded + if gen_kwargs.get("max_length") is not None and generated_tokens.shape[-1] < gen_kwargs["max_length"]: + generated_tokens = self._pad_tensors_to_max_len(generated_tokens, gen_kwargs["max_length"]) + elif gen_kwargs.get("max_new_tokens") is not None and generated_tokens.shape[-1] < ( + gen_kwargs["max_new_tokens"] + 1 + ): + generated_tokens = self._pad_tensors_to_max_len(generated_tokens, gen_kwargs["max_new_tokens"] + 1) + + loss = None + + if self.args.prediction_loss_only: + return (loss, None, None) + + if has_labels: + labels = inputs["labels"] + if gen_kwargs.get("max_length") is not None and labels.shape[-1] < gen_kwargs["max_length"]: + labels = self._pad_tensors_to_max_len(labels, gen_kwargs["max_length"]) + elif gen_kwargs.get("max_new_tokens") is not None and labels.shape[-1] < ( + gen_kwargs["max_new_tokens"] + 1 + ): + labels = self._pad_tensors_to_max_len(labels, (gen_kwargs["max_new_tokens"] + 1)) + else: + labels = None + + return (loss, generated_tokens, labels) + + def _pad_tensors_to_max_len(self, tensor, max_length): + if self.tokenizer is not None and hasattr(self.tokenizer, "pad_token_id"): + # If PAD token is not defined at least EOS token has to be defined + pad_token_id = ( + self.tokenizer.pad_token_id if self.tokenizer.pad_token_id is not None else self.tokenizer.eos_token_id + ) + else: + if self.model.config.pad_token_id is not None: + pad_token_id = self.model.config.pad_token_id + else: + raise ValueError("Pad_token_id must be set in the configuration of the model, in order to pad tensors") + + padded_tensor = pad_token_id * torch.ones( + (tensor.shape[0], max_length), dtype=tensor.dtype, device=tensor.device + ) + padded_tensor[:, : tensor.shape[-1]] = tensor + return padded_tensor diff --git a/PyTorch/built-in/foundation/CodeGeeX2/ptuning/web_demo.py b/PyTorch/built-in/foundation/CodeGeeX2/ptuning/web_demo.py new file mode 100644 index 0000000000000000000000000000000000000000..b4bb160274368dcfc04a1ea036e7df8b4673fd8d --- /dev/null +++ b/PyTorch/built-in/foundation/CodeGeeX2/ptuning/web_demo.py @@ -0,0 +1,167 @@ +import os, sys + +import gradio as gr +import mdtex2html + +import torch +import transformers +from transformers import ( + AutoConfig, + AutoModel, + AutoTokenizer, + AutoTokenizer, + DataCollatorForSeq2Seq, + HfArgumentParser, + Seq2SeqTrainingArguments, + set_seed, +) + +from arguments import ModelArguments, DataTrainingArguments + + +model = None +tokenizer = None + +"""Override Chatbot.postprocess""" + + +def postprocess(self, y): + if y is None: + return [] + for i, (message, response) in enumerate(y): + y[i] = ( + None if message is None else mdtex2html.convert((message)), + None if response is None else mdtex2html.convert(response), + ) + return y + + +gr.Chatbot.postprocess = postprocess + + +def parse_text(text): + """copy from https://github.com/GaiZhenbiao/ChuanhuChatGPT/""" + lines = text.split("\n") + lines = [line for line in lines if line != ""] + count = 0 + for i, line in enumerate(lines): + if "```" in line: + count += 1 + items = line.split('`') + if count % 2 == 1: + lines[i] = f'
'
+            else:
+                lines[i] = f'
' + else: + if i > 0: + if count % 2 == 1: + line = line.replace("`", "\`") + line = line.replace("<", "<") + line = line.replace(">", ">") + line = line.replace(" ", " ") + line = line.replace("*", "*") + line = line.replace("_", "_") + line = line.replace("-", "-") + line = line.replace(".", ".") + line = line.replace("!", "!") + line = line.replace("(", "(") + line = line.replace(")", ")") + line = line.replace("$", "$") + lines[i] = "
"+line + text = "".join(lines) + return text + + +def predict(input, chatbot, max_length, top_p, temperature, history, past_key_values): + chatbot.append((parse_text(input), "")) + for response, history, past_key_values in model.stream_chat(tokenizer, input, history, past_key_values=past_key_values, + return_past_key_values=True, + max_length=max_length, top_p=top_p, + temperature=temperature): + chatbot[-1] = (parse_text(input), parse_text(response)) + + yield chatbot, history, past_key_values + + +def reset_user_input(): + return gr.update(value='') + + +def reset_state(): + return [], [], None + + +with gr.Blocks() as demo: + gr.HTML("""

ChatGLM2-6B

""") + + chatbot = gr.Chatbot() + with gr.Row(): + with gr.Column(scale=4): + with gr.Column(scale=12): + user_input = gr.Textbox(show_label=False, placeholder="Input...", lines=10).style( + container=False) + with gr.Column(min_width=32, scale=1): + submitBtn = gr.Button("Submit", variant="primary") + with gr.Column(scale=1): + emptyBtn = gr.Button("Clear History") + max_length = gr.Slider(0, 32768, value=8192, step=1.0, label="Maximum length", interactive=True) + top_p = gr.Slider(0, 1, value=0.8, step=0.01, label="Top P", interactive=True) + temperature = gr.Slider(0, 1, value=0.95, step=0.01, label="Temperature", interactive=True) + + history = gr.State([]) + past_key_values = gr.State(None) + + submitBtn.click(predict, [user_input, chatbot, max_length, top_p, temperature, history, past_key_values], + [chatbot, history, past_key_values], show_progress=True) + submitBtn.click(reset_user_input, [], [user_input]) + + emptyBtn.click(reset_state, outputs=[chatbot, history, past_key_values], show_progress=True) + + +def main(): + global model, tokenizer + + parser = HfArgumentParser(( + ModelArguments)) + if len(sys.argv) == 2 and sys.argv[1].endswith(".json"): + # If we pass only one argument to the script and it's the path to a json file, + # let's parse it to get our arguments. + model_args = parser.parse_json_file(json_file=os.path.abspath(sys.argv[1]))[0] + else: + model_args = parser.parse_args_into_dataclasses()[0] + + tokenizer = AutoTokenizer.from_pretrained( + model_args.model_name_or_path, trust_remote_code=True) + config = AutoConfig.from_pretrained( + model_args.model_name_or_path, trust_remote_code=True) + + config.pre_seq_len = model_args.pre_seq_len + config.prefix_projection = model_args.prefix_projection + + if model_args.ptuning_checkpoint is not None: + print(f"Loading prefix_encoder weight from {model_args.ptuning_checkpoint}") + model = AutoModel.from_pretrained(model_args.model_name_or_path, config=config, trust_remote_code=True) + prefix_state_dict = torch.load(os.path.join(model_args.ptuning_checkpoint, "pytorch_model.bin")) + new_prefix_state_dict = {} + for k, v in prefix_state_dict.items(): + if k.startswith("transformer.prefix_encoder."): + new_prefix_state_dict[k[len("transformer.prefix_encoder."):]] = v + model.transformer.prefix_encoder.load_state_dict(new_prefix_state_dict) + else: + model = AutoModel.from_pretrained(model_args.model_name_or_path, config=config, trust_remote_code=True) + + if model_args.quantization_bit is not None: + print(f"Quantized to {model_args.quantization_bit} bit") + model = model.quantize(model_args.quantization_bit) + model = model.cuda() + if model_args.pre_seq_len is not None: + # P-tuning v2 + model.transformer.prefix_encoder.float() + + model = model.eval() + demo.queue().launch(share=False, inbrowser=True) + + + +if __name__ == "__main__": + main() \ No newline at end of file diff --git a/PyTorch/built-in/foundation/CodeGeeX2/ptuning/web_demo.sh b/PyTorch/built-in/foundation/CodeGeeX2/ptuning/web_demo.sh new file mode 100644 index 0000000000000000000000000000000000000000..b9465cb5362b3390a4dfea5cacdbf28c78e1e1be --- /dev/null +++ b/PyTorch/built-in/foundation/CodeGeeX2/ptuning/web_demo.sh @@ -0,0 +1,7 @@ +PRE_SEQ_LEN=128 + +CUDA_VISIBLE_DEVICES=0 python3 web_demo.py \ + --model_name_or_path THUDM/chatglm2-6b \ + --ptuning_checkpoint output/adgen-chatglm2-6b-pt-128-2e-2/checkpoint-3000 \ + --pre_seq_len $PRE_SEQ_LEN + diff --git a/PyTorch/built-in/foundation/CodeGeeX2/requirements.txt b/PyTorch/built-in/foundation/CodeGeeX2/requirements.txt new file mode 100644 index 0000000000000000000000000000000000000000..2f29fe50fa32f3bc6a82a56372d66b5d1585a13d --- /dev/null +++ b/PyTorch/built-in/foundation/CodeGeeX2/requirements.txt @@ -0,0 +1,7 @@ +protobuf +transformers>=4.30.2 +accelerate +cpm_kernels +torch>=2.0 +sentencepiece +gradio \ No newline at end of file diff --git a/PyTorch/built-in/foundation/CodeGeeX2/resources/wechat.md b/PyTorch/built-in/foundation/CodeGeeX2/resources/wechat.md new file mode 100644 index 0000000000000000000000000000000000000000..8fbc090709455f65ff59b4f80796ce61b854a8f0 --- /dev/null +++ b/PyTorch/built-in/foundation/CodeGeeX2/resources/wechat.md @@ -0,0 +1,6 @@ +
+ + +

扫码关注公众号加入「CodeGeeX交流群」

+

Scan the QR code to join the "CodeGeeX WeChat Group"

+
diff --git a/PyTorch/built-in/foundation/CodeGeeX2/scripts/hostlist b/PyTorch/built-in/foundation/CodeGeeX2/scripts/hostlist new file mode 100644 index 0000000000000000000000000000000000000000..a9eef97c3b833092b40fb74fd8f1984b1fb52ee3 --- /dev/null +++ b/PyTorch/built-in/foundation/CodeGeeX2/scripts/hostlist @@ -0,0 +1 @@ +172.0.0.1 diff --git a/PyTorch/built-in/foundation/CodeGeeX2/scripts/run_humanevalx.sh b/PyTorch/built-in/foundation/CodeGeeX2/scripts/run_humanevalx.sh new file mode 100644 index 0000000000000000000000000000000000000000..8285037aa37a2de9bf4c680944540b43073af892 --- /dev/null +++ b/PyTorch/built-in/foundation/CodeGeeX2/scripts/run_humanevalx.sh @@ -0,0 +1,171 @@ +#!/bin/bash +# This script is used to generate solutions of HumanEval-X. + +# Examples (MODE=(gen, eval, both)): +# MODE=gen bash ./scripts/run_humanevalx.sh + +if [ -z "$MODE" ] +then + MODE="both" +fi + +SCRIPT_PATH=$(realpath "$0") +SCRIPT_DIR=$(dirname "$SCRIPT_PATH") +MAIN_DIR=$(dirname "$SCRIPT_DIR") + +# enviroment settings +HOSTLIST=$SCRIPT_DIR/hostlist +WORLD_SIZE=1 +DATASET=humanevalx +GENERATION_MODE=completion +MODEL_NAME=codegeex2-6b +MODEL_PATH=../ptuning/output/codegeex2-6b-ft-2e-5/checkpoint-100 +N_CPU_WORKERS=16 +TIMEOUT=5 + +# generation settings +## pass@1 greedy +NUM_SAMPLES=1 +MICRO_BSZ=1 +TEMP=1.0 +TOPK=1 +TOPP=1.0 +MAX_LENGTH=1024 +SEED=42 +GREEDY=1 + +## pass@1 estimated +# NUM_SAMPLES=20 +# MICRO_BSZ=1 +# TEMP=0.2 +# TOPK=0 +# TOPP=0.95 +# MAX_LENGTH=1024 +# SEED=42 +# GREEDY=0 + +## pass@10 & pass@100 +# NUM_SAMPLES=200 +# MICRO_BSZ=4 +# TEMP=0.8 +# TOPK=0 +# TOPP=0.95 +# MAX_LENGTH=1024 +# SEED=42 +# GREEDY=0 + +for l in python java js cpp go rust; +do + LANGUAGE=$l + DATA_DIR=$MAIN_DIR/benchmark/$DATASET/ + DATA_PATH=$DATA_DIR/$DATASET\_$LANGUAGE.jsonl.gz + OUTPUT_PATH=$MAIN_DIR/output/$DATASET/$LANGUAGE + TODAY=$(date +%y%m%d) + CHANNEL_PORT=$(expr $RANDOM + 5000) + MASTER_PORT=$(expr $RANDOM + 8000) + JOB_ID=$MODEL_NAME-$LANGUAGE-greedy$GREEDY-ns$NUM_SAMPLES-t$TEMP-topp$TOPP-seed$SEED + mkdir -p "$OUTPUT_PATH/$JOB_ID" + + # evaluation settings + EVAL_INPUT_PATH=$OUTPUT_PATH/$JOB_ID + EVAL_OUTPUT_PATH=$OUTPUT_PATH/$JOB_ID + + # nccl options + OPTIONS_NCCL="export NCCL_DEBUG=warn; export NCCL_IB_DISABLE=0; export NCCL_IB_GID_INDEX=3" + OPTIONS_PATH="export PATH=$PATH; export LD_LIBRARY_PATH=$LD_LIBRARY_PATH" + CWD=$(pwd) + + gen_func() { + echo "Generating......" + # set master ip for zmq server + if [ -z "$HOSTLIST" ]; then + ZMQ_ADDR=$(hostname -i) + echo "$ZMQ_ADDR" > "./hostfile" + HOSTLIST="./hostfile" + else + ZMQ_ADDR=$(cat $HOSTLIST | head -n 1) + fi + echo "master_ip: $ZMQ_ADDR" + + # run generation + RUN_CMD="python \ + $MAIN_DIR/evaluation/generation.py \ + --hostfile $HOSTLIST \ + --channel-ip $ZMQ_ADDR \ + --channel-port $CHANNEL_PORT \ + --master-port $MASTER_PORT \ + --model-path $MODEL_PATH \ + --temperature $TEMP \ + --top-p $TOPP \ + --top-k $TOPK \ + --greedy $GREEDY \ + --max-length $MAX_LENGTH \ + --micro-batch-size $MICRO_BSZ \ + --samples-per-problem $NUM_SAMPLES \ + --model-name $MODEL_NAME \ + --dataset-type $DATASET \ + --language-type $LANGUAGE \ + --generation-mode $GENERATION_MODE \ + --data-path $DATA_PATH \ + --output-path $OUTPUT_PATH/$JOB_ID \ + --log-path $OUTPUT_PATH/$JOB_ID/$TODAY-generation.log \ + --gen-node-world-size $WORLD_SIZE \ + --seed $SEED" + + RUN_CMD="$OPTIONS_NCCL; $OPTIONS_PATH; $RUN_CMD" + RUN_CMD="cd $CWD; $RUN_CMD" + + if (( WORLD_SIZE != 1 )); then + RUN_CMD="pdsh -R ssh -w ^$HOSTLIST \"$RUN_CMD\"" + fi + + eval "$RUN_CMD" + } + + eval_func() { + echo "Evaluating......" + + if [ $LANGUAGE = rust ]; then + TIMEOUT=300 + echo "Setting timeout to $TIMEOUT for Rust" + fi + RUN_CMD="python \ + $MAIN_DIR/evaluation/evaluation.py \ + --input_path $EVAL_INPUT_PATH \ + --output_path $EVAL_OUTPUT_PATH \ + --log-path $OUTPUT_PATH/$JOB_ID/$TODAY-evaluation.log \ + --model_name $MODEL_NAME \ + --language_type $LANGUAGE \ + --dataset_type $DATASET \ + --generation_mode $GENERATION_MODE \ + --n_workers $N_CPU_WORKERS \ + --tmp_dir $MAIN_DIR/benchmark/$DATASET/$LANGUAGE \ + --problem_file $DATA_PATH \ + --timeout $TIMEOUT" + + # inspecting results + INSPECT_CMD="python \ + $MAIN_DIR/evaluation/inspect_jsonl.py \ + --data_path $EVAL_OUTPUT_PATH/result-$JOB_ID.jsonl \ + --log-path $OUTPUT_PATH/$JOB_ID/$TODAY-inspect.txt" + + eval "$RUN_CMD && $INSPECT_CMD" + } + + case $MODE in + "gen") + gen_func + ;; + "eval") + eval_func + ;; + "both") + gen_func + eval_func + ;; + *) + echo "Unsupported MODE (gen, eval, both): $MODE" + exit 1 + ;; + esac +done diff --git a/PyTorch/built-in/foundation/CodeGeeX2/scripts/sanity_check.sh b/PyTorch/built-in/foundation/CodeGeeX2/scripts/sanity_check.sh new file mode 100644 index 0000000000000000000000000000000000000000..751435e3581c72aae7e0a7e6fc9889e75b48a2bb --- /dev/null +++ b/PyTorch/built-in/foundation/CodeGeeX2/scripts/sanity_check.sh @@ -0,0 +1,51 @@ +#!/bin/bash +# This script is used to check the correctness of code generation benchmarks. + +SCRIPT_PATH=$(realpath "$0") +SCRIPT_DIR=$(dirname "$SCRIPT_PATH") +MAIN_DIR=$(dirname "$SCRIPT_DIR") + +# enviroment settings +DATASET=humanevalx +GENERATION_MODE=completion +N_CPU_WORKERS=16 +TIMEOUT=5 + +# Check HumanEval-X +for l in python java js cpp go rust; +do + LANGUAGE=$l + echo "Evaluating $l" + DATA_DIR=$MAIN_DIR/benchmark/$DATASET/ + DATA_PATH=$DATA_DIR/$DATASET\_$LANGUAGE.jsonl.gz + OUTPUT_PATH=$MAIN_DIR/output/$DATASET/$LANGUAGE + + JOB_ID=sanity-check-$LANGUAGE + mkdir -p "$OUTPUT_PATH/$JOB_ID" + + # evaluation settings + EVAL_INPUT_PATH=$DATA_PATH + EVAL_OUTPUT_PATH=$OUTPUT_PATH/$JOB_ID + + if [ $LANGUAGE = rust ]; then + TIMEOUT=300 + echo "Setting timeout to $TIMEOUT for Rust" + fi + + RUN_CMD="python \ + $MAIN_DIR/evaluation/evaluation.py \ + --test_groundtruth=True \ + --input_path $EVAL_INPUT_PATH \ + --output_path $EVAL_OUTPUT_PATH \ + --log-path $OUTPUT_PATH/$JOB_ID/$TODAY-evaluation.log \ + --model_name $MODEL_NAME \ + --language_type $LANGUAGE \ + --dataset_type $DATASET \ + --generation_mode $GENERATION_MODE \ + --n_workers $N_CPU_WORKERS \ + --tmp_dir $MAIN_DIR/benchmark/$DATASET/$LANGUAGE \ + --problem_file $DATA_PATH \ + --timeout $TIMEOUT" + + eval "$RUN_CMD" +done diff --git a/PyTorch/built-in/foundation/CodeGeeX2/utils.py b/PyTorch/built-in/foundation/CodeGeeX2/utils.py new file mode 100644 index 0000000000000000000000000000000000000000..78fa5e8a44441c99927b9a757f13815857cc24c8 --- /dev/null +++ b/PyTorch/built-in/foundation/CodeGeeX2/utils.py @@ -0,0 +1,59 @@ +import os +from typing import Dict, Tuple, Union, Optional + +from torch.nn import Module +from transformers import AutoModel + + +def auto_configure_device_map(num_gpus: int) -> Dict[str, int]: + # transformer.word_embeddings 占用1层 + # transformer.final_layernorm 和 lm_head 占用1层 + # transformer.layers 占用 28 层 + # 总共30层分配到num_gpus张卡上 + num_trans_layers = 28 + per_gpu_layers = 30 / num_gpus + + # bugfix: 在linux中调用torch.embedding传入的weight,input不在同一device上,导致RuntimeError + # windows下 model.device 会被设置成 transformer.word_embeddings.device + # linux下 model.device 会被设置成 lm_head.device + # 在调用chat或者stream_chat时,input_ids会被放到model.device上 + # 如果transformer.word_embeddings.device和model.device不同,则会导致RuntimeError + # 因此这里将transformer.word_embeddings,transformer.final_layernorm,lm_head都放到第一张卡上 + # 本文件来源于https://github.com/THUDM/ChatGLM-6B/blob/main/utils.py + # 仅此处做少许修改以支持ChatGLM2 + device_map = { + 'transformer.embedding.word_embeddings': 0, + 'transformer.encoder.final_layernorm': 0, + 'transformer.output_layer': 0, + 'transformer.rotary_pos_emb': 0, + 'lm_head': 0 + } + + used = 2 + gpu_target = 0 + for i in range(num_trans_layers): + if used >= per_gpu_layers: + gpu_target += 1 + used = 0 + assert gpu_target < num_gpus + device_map[f'transformer.encoder.layers.{i}'] = gpu_target + used += 1 + + return device_map + + +def load_model_on_gpus(checkpoint_path: Union[str, os.PathLike], num_gpus: int = 2, + device_map: Optional[Dict[str, int]] = None, **kwargs) -> Module: + if num_gpus < 2 and device_map is None: + model = AutoModel.from_pretrained(checkpoint_path, trust_remote_code=True, **kwargs).half().cuda() + else: + from accelerate import dispatch_model + + model = AutoModel.from_pretrained(checkpoint_path, trust_remote_code=True, **kwargs).half() + + if device_map is None: + device_map = auto_configure_device_map(num_gpus) + + model = dispatch_model(model, device_map=device_map) + + return model diff --git a/PyTorch/built-in/foundation/CodeGeeX2/web_demo.py b/PyTorch/built-in/foundation/CodeGeeX2/web_demo.py new file mode 100644 index 0000000000000000000000000000000000000000..1af24c962943c74199712269e45a94a02a4b2cc5 --- /dev/null +++ b/PyTorch/built-in/foundation/CodeGeeX2/web_demo.py @@ -0,0 +1,108 @@ +from transformers import AutoModel, AutoTokenizer +import gradio as gr +import mdtex2html +from utils import load_model_on_gpus + +tokenizer = AutoTokenizer.from_pretrained("THUDM/chatglm2-6b", trust_remote_code=True) +model = AutoModel.from_pretrained("THUDM/chatglm2-6b", trust_remote_code=True).cuda() +# 多显卡支持，使用下面两行代替上面一行，将num_gpus改为你实际的显卡数量 +# from utils import load_model_on_gpus +# model = load_model_on_gpus("THUDM/chatglm2-6b", num_gpus=2) +model = model.eval() + +"""Override Chatbot.postprocess""" + + +def postprocess(self, y): + if y is None: + return [] + for i, (message, response) in enumerate(y): + y[i] = ( + None if message is None else mdtex2html.convert((message)), + None if response is None else mdtex2html.convert(response), + ) + return y + + +gr.Chatbot.postprocess = postprocess + + +def parse_text(text): + """copy from https://github.com/GaiZhenbiao/ChuanhuChatGPT/""" + lines = text.split("\n") + lines = [line for line in lines if line != ""] + count = 0 + for i, line in enumerate(lines): + if "```" in line: + count += 1 + items = line.split('`') + if count % 2 == 1: + lines[i] = f'
'
+            else:
+                lines[i] = f'
' + else: + if i > 0: + if count % 2 == 1: + line = line.replace("`", "\`") + line = line.replace("<", "<") + line = line.replace(">", ">") + line = line.replace(" ", " ") + line = line.replace("*", "*") + line = line.replace("_", "_") + line = line.replace("-", "-") + line = line.replace(".", ".") + line = line.replace("!", "!") + line = line.replace("(", "(") + line = line.replace(")", ")") + line = line.replace("$", "$") + lines[i] = "
"+line + text = "".join(lines) + return text + + +def predict(input, chatbot, max_length, top_p, temperature, history, past_key_values): + chatbot.append((parse_text(input), "")) + for response, history, past_key_values in model.stream_chat(tokenizer, input, history, past_key_values=past_key_values, + return_past_key_values=True, + max_length=max_length, top_p=top_p, + temperature=temperature): + chatbot[-1] = (parse_text(input), parse_text(response)) + + yield chatbot, history, past_key_values + + +def reset_user_input(): + return gr.update(value='') + + +def reset_state(): + return [], [], None + + +with gr.Blocks() as demo: + gr.HTML("""

ChatGLM2-6B

""") + + chatbot = gr.Chatbot() + with gr.Row(): + with gr.Column(scale=4): + with gr.Column(scale=12): + user_input = gr.Textbox(show_label=False, placeholder="Input...", lines=10).style( + container=False) + with gr.Column(min_width=32, scale=1): + submitBtn = gr.Button("Submit", variant="primary") + with gr.Column(scale=1): + emptyBtn = gr.Button("Clear History") + max_length = gr.Slider(0, 32768, value=8192, step=1.0, label="Maximum length", interactive=True) + top_p = gr.Slider(0, 1, value=0.8, step=0.01, label="Top P", interactive=True) + temperature = gr.Slider(0, 1, value=0.95, step=0.01, label="Temperature", interactive=True) + + history = gr.State([]) + past_key_values = gr.State(None) + + submitBtn.click(predict, [user_input, chatbot, max_length, top_p, temperature, history, past_key_values], + [chatbot, history, past_key_values], show_progress=True) + submitBtn.click(reset_user_input, [], [user_input]) + + emptyBtn.click(reset_state, outputs=[chatbot, history, past_key_values], show_progress=True) + +demo.queue().launch(share=False, inbrowser=True) diff --git a/PyTorch/built-in/foundation/CodeGeeX2/web_demo2.py b/PyTorch/built-in/foundation/CodeGeeX2/web_demo2.py new file mode 100644 index 0000000000000000000000000000000000000000..203cbdc7c8d6d212a26bd63b4770781abd5512f1 --- /dev/null +++ b/PyTorch/built-in/foundation/CodeGeeX2/web_demo2.py @@ -0,0 +1,70 @@ +from transformers import AutoModel, AutoTokenizer +import streamlit as st + + +st.set_page_config( + page_title="ChatGLM2-6b 演示", + page_icon=":robot:", + layout='wide' +) + + +@st.cache_resource +def get_model(): + tokenizer = AutoTokenizer.from_pretrained("THUDM/chatglm2-6b", trust_remote_code=True) + model = AutoModel.from_pretrained("THUDM/chatglm2-6b", trust_remote_code=True).cuda() + # 多显卡支持，使用下面两行代替上面一行，将num_gpus改为你实际的显卡数量 + # from utils import load_model_on_gpus + # model = load_model_on_gpus("THUDM/chatglm2-6b", num_gpus=2) + model = model.eval() + return tokenizer, model + + +tokenizer, model = get_model() + +st.title("ChatGLM2-6B") + +max_length = st.sidebar.slider( + 'max_length', 0, 32768, 8192, step=1 +) +top_p = st.sidebar.slider( + 'top_p', 0.0, 1.0, 0.8, step=0.01 +) +temperature = st.sidebar.slider( + 'temperature', 0.0, 1.0, 0.8, step=0.01 +) + +if 'history' not in st.session_state: + st.session_state.history = [] + +if 'past_key_values' not in st.session_state: + st.session_state.past_key_values = None + +for i, (query, response) in enumerate(st.session_state.history): + with st.chat_message(name="user", avatar="user"): + st.markdown(query) + with st.chat_message(name="assistant", avatar="assistant"): + st.markdown(response) +with st.chat_message(name="user", avatar="user"): + input_placeholder = st.empty() +with st.chat_message(name="assistant", avatar="assistant"): + message_placeholder = st.empty() + +prompt_text = st.text_area(label="用户命令输入", + height=100, + placeholder="请在这儿输入您的命令") + +button = st.button("发送", key="predict") + +if button: + input_placeholder.markdown(prompt_text) + history, past_key_values = st.session_state.history, st.session_state.past_key_values + for response, history, past_key_values in model.stream_chat(tokenizer, prompt_text, history, + past_key_values=past_key_values, + max_length=max_length, top_p=top_p, + temperature=temperature, + return_past_key_values=True): + message_placeholder.markdown(response) + + st.session_state.history = history + st.session_state.past_key_values = past_key_values