upload Hy-MT2 files

.gitattributes

@@ -34,3 +34,5 @@ saved_model/**/* filter=lfs diff=lfs merge=lfs -text
 *.zst filter=lfs diff=lfs merge=lfs -text
 *tfevents* filter=lfs diff=lfs merge=lfs -text
 tokenizer.json filter=lfs diff=lfs merge=lfs -text
+HY_MT2_0_Technical_Report.pdf filter=lfs diff=lfs merge=lfs -text
+imgs/main_result.png filter=lfs diff=lfs merge=lfs -text

HY_MT2_0_Technical_Report.pdf

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:ce04272779bf9df1e47ba64132008ff95d9416c4111d266927390749812d1b72
+size 2597469

LICENSE.txt

@@ -0,0 +1,80 @@
+TENCENT HY COMMUNITY LICENSE AGREEMENT
+Tencent Hy-MT2 Release Date: May 21, 2026
+THIS LICENSE AGREEMENT DOES NOT APPLY IN THE EUROPEAN UNION, AND IS EXPRESSLY LIMITED TO THE TERRITORY, AS DEFINED BELOW.
+
+By clicking to agree or by using, reproducing, modifying, distributing, performing or displaying any portion or element of the Tencent HY Works, including via any Hosted Service, You will be deemed to have recognized and accepted the content of this Agreement, which is effective immediately.
+
+1.	DEFINITIONS.
+a.	“Acceptable Use Policy” shall mean the policy made available by Tencent as set forth in the Exhibit A.
+b.	“Agreement” shall mean the terms and conditions for use, reproduction, distribution, modification, performance and displaying of Tencent HY Works or any portion or element thereof set forth herein.
+c.	“Documentation” shall mean the specifications, manuals and documentation for Tencent HY made publicly available by Tencent.
+d.	“Hosted Service” shall mean a hosted service offered via an application programming interface (API), web access, or any other electronic or remote means.
+e.	“Licensee,” “You” or “Your” shall mean a natural person or legal entity exercising the rights granted by this Agreement and/or using the Tencent HY Works for any purpose and in any field of use.
+f.	“Materials” shall mean, collectively, Tencent’s proprietary Tencent HY and Documentation (and any portion thereof) as made available by Tencent under this Agreement.
+g.	“Model Derivatives” shall mean all: (i) modifications to Tencent HY or any Model Derivative of Tencent HY; (ii) works based on Tencent HY or any Model Derivative of Tencent HY; or (iii) any other machine learning model which is created by transfer of patterns of the weights, parameters, operations, or Output of Tencent HY or any Model Derivative of Tencent HY, to that model in order to cause that model to perform similarly to Tencent HY or a Model Derivative of Tencent HY, including distillation methods, methods that use intermediate data representations, or methods based on the generation of synthetic data Outputs by Tencent HY or a Model Derivative of Tencent HY for training that model. For clarity, Outputs by themselves are not deemed Model Derivatives.
+h.	“Output” shall mean the information and/or content output of Tencent HY or a Model Derivative that results from operating or otherwise using Tencent HY or a Model Derivative, including via a Hosted Service.
+i.	“Tencent,” “We” or “Us” shall mean the applicable entity or entities in the Tencent corporate family that own(s) intellectual property or other rights embodied in or utilized by the Materials.
+j.	“Tencent HY” shall mean the large language models, text/image/video/audio/3D generation models, and multimodal large language models and their software and algorithms, including trained model weights, parameters (including optimizer states), machine-learning model code, inference-enabling code, training-enabling code, fine-tuning enabling code and other elements of the foregoing made publicly available by Us, including, without limitation to, Tencent Hy-MT2-1.8B released at https://huggingface.co/tencent/Hy-MT2-1.8B, https://modelscope.cn/models/Tencent-Hunyuan/Hy-MT2-1.8B; Tencent Hy-MT2-7B released at https://huggingface.co/tencent/Hy-MT2-7B, https://modelscope.cn/models/Tencent-Hunyuan/Hy-MT2-7B; Tencent Hy-MT2-30B-A3B released at https://huggingface.co/tencent/Hy-MT2-30B-A3B, https://modelscope.cn/models/Tencent-Hunyuan/Hy-MT2-30B-A3B; Tencent Hy-MT2-1.8B-FP8 released at https://huggingface.co/tencent/Hy-MT2-1.8B-FP8, https://modelscope.cn/models/Tencent-Hunyuan/Hy-MT2-1.8B-FP8; Tencent Hy-MT2-7B-FP8 released at https://huggingface.co/tencent/Hy-MT2-7B-FP8, https://modelscope.cn/models/Tencent-Hunyuan/Hy-MT2-7B-FP8; Tencent Hy-MT2-30B-A3B-FP8 released at https://huggingface.co/tencent/Hy-MT2-30B-A3B-FP8, https://modelscope.cn/models/Tencent-Hunyuan/Hy-MT2-30B-A3B-FP8; Hy-MT2-1.8B-GGUF released at https://huggingface.co/tencent/Hy-MT2-1.8B-GGUF, https://modelscope.cn/models/Tencent-Hunyuan/Hy-MT2-1.8B-GGUF; Hy-MT2-7B-GGUF released at https://huggingface.co/tencent/Hy-MT2-7B-GGUF, https://modelscope.cn/models/Tencent-Hunyuan/Hy-MT2-7B-GGUF. 
+k.	“Tencent HY Works” shall mean: (i) the Materials; (ii) Model Derivatives; and (iii) all derivative works thereof.
+l.	“Territory” shall mean the worldwide territory, excluding the territory of the European Union. 
+m.	“Third Party” or “Third Parties” shall mean individuals or legal entities that are not under common control with Us or You.
+n.	“including” shall mean including but not limited to.
+2.	GRANT OF RIGHTS.
+We grant You, for the Territory only, a non-exclusive, non-transferable and royalty-free limited license under Tencent’s intellectual property or other rights owned by Us embodied in or utilized by the Materials to use, reproduce, distribute, create derivative works of (including Model Derivatives), and make modifications to the Materials, only in accordance with the terms of this Agreement and the Acceptable Use Policy, and You must not violate (or encourage or permit anyone else to violate) any term of this Agreement or the Acceptable Use Policy.
+3.	DISTRIBUTION.
+You may, subject to Your compliance with this Agreement, distribute or make available to Third Parties the Tencent HY Works, exclusively in the Territory, provided that You meet all of the following conditions:
+a.	You must provide all such Third Party recipients of the Tencent HY Works or products or services using them a copy of this Agreement;
+b.	You must cause any modified files to carry prominent notices stating that You changed the files;
+c.	You are encouraged to: (i) publish at least one technology introduction blogpost or one public statement expressing Your experience of using the Tencent HY Works; and (ii) mark the products or services developed by using the Tencent HY Works to indicate that the product/service is “Powered by Tencent HY”; and
+d.	All distributions to Third Parties (other than through a Hosted Service) must be accompanied by a “Notice” text file that contains the following notice: “Tencent HY is licensed under the Tencent HY Community License Agreement, Copyright © 2026 Tencent. All Rights Reserved. The trademark rights of “Tencent HY” are owned by Tencent or its affiliate.”
+e.	In the event that You use, integrate, implement, or otherwise deploy the Tencent HY Works, in whole or in part, to provide, enable, or support any service, product, or functionality to third parties, You shall clearly, accurately, and prominently disclose to all end users the full legal name and entity of the actual provider of such service, product, or functionality. You shall expressly and conspicuously state that Tencent is not affiliated with, associated with, sponsoring, or endorsing any such service, product, or functionality. You shall not use or display any name, logo, trademark, trade name, or other indicia of Tencent in any manner that could be construed as, or be likely to create, confusion, deception, or a false impression regarding any relationship, affiliation, sponsorship, or endorsement by Tencent.
+You may add Your own copyright statement to Your modifications and, except as set forth in this Section and in Section 5, may provide additional or different license terms and conditions for use, reproduction, or distribution of Your modifications, or for any such Model Derivatives as a whole, provided Your use, reproduction, modification, distribution, performance and display of the work otherwise complies with the terms and conditions of this Agreement (including as regards the Territory). If You receive Tencent HY Works from a Licensee as part of an integrated end user product, then this Section 3 of this Agreement will not apply to You.
+4.	ADDITIONAL COMMERCIAL TERMS.
+If, on the Tencent HY version release date, the monthly active users of all products or services made available by or for Licensee is greater than 100 million monthly active users in the preceding calendar month, You must request a license from Tencent, which Tencent may grant to You in its sole discretion, and You are not authorized to exercise any of the rights under this Agreement unless or until Tencent otherwise expressly grants You such rights.
+5.	RULES OF USE.
+a.	Your use of the Tencent HY Works must comply with applicable laws and regulations (including trade compliance laws and regulations) and adhere to the Acceptable Use Policy for the Tencent HY Works, which is hereby incorporated by reference into this Agreement. You must include the use restrictions referenced in these Sections 5(a) and 5(b) as an enforceable provision in any agreement (e.g., license agreement, terms of use, etc.) governing the use and/or distribution of Tencent HY Works and You must provide notice to subsequent users to whom You distribute that Tencent HY Works are subject to the use restrictions in these Sections 5(a) and 5(b).
+b.	You must not use the Tencent HY Works or any Output or results of the Tencent HY Works to improve any other AI model (other than Tencent HY or Model Derivatives thereof).
+c.	You must not use, reproduce, modify, distribute, or display the Tencent HY Works, Output or results of the Tencent HY Works outside the Territory. Any such use outside the Territory is unlicensed and unauthorized under this Agreement.
+6.	INTELLECTUAL PROPERTY.
+a.	Subject to Tencent’s ownership of Tencent HY Works made by or for Tencent and intellectual property rights therein, conditioned upon Your compliance with the terms and conditions of this Agreement, as between You and Tencent, You will be the owner of any derivative works and modifications of the Materials and any Model Derivatives that are made by or for You.
+b.	No trademark licenses are granted under this Agreement, and in connection with the Tencent HY Works, Licensee may not use any name or mark owned by or associated with Tencent or any of its affiliates, except as required for reasonable and customary use in describing and distributing the Tencent HY Works. Tencent hereby grants You a license to use “Tencent HY” (the “Mark”) in the Territory solely as required to comply with the provisions of Section 3(c), provided that You comply with any applicable laws related to trademark protection. All goodwill arising out of Your use of the Mark will inure to the benefit of Tencent.
+c.	If You commence a lawsuit or other proceedings (including a cross-claim or counterclaim in a lawsuit) against Us or any person or entity alleging that the Materials or any Output, or any portion of any of the foregoing, infringe any intellectual property or other right owned or licensable by You, then all licenses granted to You under this Agreement shall terminate as of the date such lawsuit or other proceeding is filed. You will defend, indemnify and hold harmless Us from and against any claim by any Third Party arising out of or related to Your or the Third Party’s use or distribution of the Tencent HY Works.
+d.	Tencent claims no rights in Outputs You generate. You and Your users are solely responsible for Outputs and their subsequent uses.
+7.	DISCLAIMERS OF WARRANTY AND LIMITATIONS OF LIABILITY.
+a.	We are not obligated to support, update, provide training for, or develop any further version of the Tencent HY Works or to grant any license thereto.
+b.	UNLESS AND ONLY TO THE EXTENT REQUIRED BY APPLICABLE LAW, THE TENCENT HY WORKS AND ANY OUTPUT AND RESULTS THEREFROM ARE PROVIDED “AS IS” WITHOUT ANY EXPRESS OR IMPLIED WARRANTIES OF ANY KIND INCLUDING ANY WARRANTIES OF TITLE, MERCHANTABILITY, NONINFRINGEMENT, COURSE OF DEALING, USAGE OF TRADE, OR FITNESS FOR A PARTICULAR PURPOSE. YOU ARE SOLELY RESPONSIBLE FOR DETERMINING THE APPROPRIATENESS OF USING, REPRODUCING, MODIFYING, PERFORMING, DISPLAYING OR DISTRIBUTING ANY OF THE TENCENT HY WORKS OR OUTPUTS AND ASSUME ANY AND ALL RISKS ASSOCIATED WITH YOUR OR A THIRD PARTY’S USE OR DISTRIBUTION OF ANY OF THE TENCENT HY WORKS OR OUTPUTS AND YOUR EXERCISE OF RIGHTS AND PERMISSIONS UNDER THIS AGREEMENT.
+c.	TO THE FULLEST EXTENT PERMITTED BY APPLICABLE LAW, IN NO EVENT SHALL TENCENT OR ITS AFFILIATES BE LIABLE UNDER ANY THEORY OF LIABILITY, WHETHER IN CONTRACT, TORT, NEGLIGENCE, PRODUCTS LIABILITY, OR OTHERWISE, FOR ANY DAMAGES, INCLUDING ANY DIRECT, INDIRECT, SPECIAL, INCIDENTAL, EXEMPLARY, CONSEQUENTIAL OR PUNITIVE DAMAGES, OR LOST PROFITS OF ANY KIND ARISING FROM THIS AGREEMENT OR RELATED TO ANY OF THE TENCENT HY WORKS OR OUTPUTS, EVEN IF TENCENT OR ITS AFFILIATES HAVE BEEN ADVISED OF THE POSSIBILITY OF ANY OF THE FOREGOING.
+8.	SURVIVAL AND TERMINATION.
+a.	The term of this Agreement shall commence upon Your acceptance of this Agreement or access to the Materials and will continue in full force and effect until terminated in accordance with the terms and conditions herein.
+b.	We may terminate this Agreement if You breach any of the terms or conditions of this Agreement. Upon termination of this Agreement, You must promptly delete and cease use of the Tencent HY Works. Sections 6(a), 6(c), 7 and 9 shall survive the termination of this Agreement.
+9.	GOVERNING LAW AND JURISDICTION.
+a.	This Agreement and any dispute arising out of or relating to it will be governed by the laws of the Hong Kong Special Administrative Region of the People’s Republic of China, without regard to conflict of law principles, and the UN Convention on Contracts for the International Sale of Goods does not apply to this Agreement.
+b.	Exclusive jurisdiction and venue for any dispute arising out of or relating to this Agreement will be a court of competent jurisdiction in the Hong Kong Special Administrative Region of the People’s Republic of China, and Tencent and Licensee consent to the exclusive jurisdiction of such court with respect to any such dispute.
+ 
+EXHIBIT A
+ACCEPTABLE USE POLICY
+
+Tencent reserves the right to update this Acceptable Use Policy from time to time.
+Last modified: December 30, 2025
+
+Tencent endeavors to promote safe and fair use of its tools and features, including Tencent HY. You agree not to use Tencent HY or Model Derivatives:
+1.	Outside the Territory;
+2.	In any way that violates any applicable national, federal, state, local, international or any other law or regulation;
+3.	To harm Yourself or others;
+4.	To repurpose or distribute output from Tencent HY or any Model Derivatives to harm Yourself or others; 
+5.	To override or circumvent the safety guardrails and safeguards We have put in place;
+6.	For the purpose of exploiting, harming or attempting to exploit or harm minors in any way;
+7.	To generate or disseminate verifiably false information and/or content with the purpose of harming others or influencing elections;
+8.	To generate or facilitate false online engagement, including fake reviews and other means of fake online engagement;
+9.	To intentionally defame, disparage or otherwise harass others;
+10.	To generate and/or disseminate malware (including ransomware) or any other content to be used for the purpose of harming electronic systems;
+11.	To generate or disseminate personal identifiable information with the purpose of harming others;
+12.	To generate or disseminate information (including images, code, posts, articles), and place the information in any public context (including –through the use of bot generated tweets), without expressly and conspicuously identifying that the information and/or content is machine generated;
+13.	To impersonate another individual without consent, authorization, or legal right;
+14.	To make high-stakes automated decisions in domains that affect an individual’s safety, rights or wellbeing (e.g., law enforcement, migration, medicine/health, management of critical infrastructure, safety components of products, essential services, credit, employment, housing, education, social scoring, or insurance);
+15.	In a manner that violates or disrespects the social ethics and moral standards of other countries or regions;
+16.	To perform, facilitate, threaten, incite, plan, promote or encourage violent extremism or terrorism;
+17.	For any use intended to discriminate against or harm individuals or groups based on protected characteristics or categories, online or offline social behavior or known or predicted personal or personality characteristics;
+18.	To intentionally exploit any of the vulnerabilities of a specific group of persons based on their age, social, physical or mental characteristics, in order to materially distort the behavior of a person pertaining to that group in a manner that causes or is likely to cause that person or another person physical or psychological harm;
+19.	For military purposes;
+20.	To engage in the unauthorized or unlicensed practice of any profession including, but not limited to, financial, legal, medical/health, or other professional practices.

README.md

@@ -0,0 +1,294 @@
+<p align="left">
+   English&nbsp;｜&nbsp;<a href="README_CN.md">中文</a>
+</p>
+<br>
+
+<p align="center">
+ <img src="imgs/logo-en.png" width="400"/> <br> 
+</p>
+
+<div align="center" style="line-height: 1;">
+
+
+[![HuggingFace](https://img.shields.io/badge/%F0%9F%A4%97%20Hugging%20Face-Tencent%20Hy-ffc107?color=ffc107&logoColor=white)](https://huggingface.co/collections/tencent/hy-mt2)
+&nbsp;&nbsp;
+[![ModelScope](https://img.shields.io/badge/ModelScope-Tencent%20Hy-624aff)](https://modelscope.cn/collections/Tencent-Hunyuan/Hy-MT2)
+
+</div>
+
+<p align="center">
+    🖥️&nbsp;<a href="https://aistudio.tencent.com/"><b>Official Website</b></a>&nbsp;&nbsp;|&nbsp;&nbsp;
+    💬&nbsp;<a href="https://github.com/Tencent-Hunyuan/Hy-MT2"><b>GitHub</b></a>&nbsp;&nbsp;|&nbsp;&nbsp;
+    🪡&nbsp;<a href="https://github.com/Tencent/AngelSlim/tree/main"><b>AngelSlim</b></a></p>
+
+## Model Introduction
+
+
+**Hy-MT2** is a multilingual machine translation model series covering both Dense and MoE architectures. It includes three fast-thinking models: **Hy-MT2-1.8B, 7B, and 30B-A3B**. The series supports translation among 33 languages and 5 ethnic minority languages / Chinese dialects, as well as multilingual instruction following. The series also provides **1.25-bit extreme quantized versions** based on AngelSlim. Among them, the 1.8B model requires only 440 MB of storage and runs 1.5x faster than traditional 4-bit inference on the Apple A15 chip.
+
+Evaluation results show that Hy-MT2 performs strongly across multiple scenarios:
+
+* **General Translation (FLORES-200)**: The average performance of the three models reaches 89.9%, 97.9%, and 98.6% of **Gemini 3.1 Pro (Think)**, respectively. Among them, the 7B and A3B models outperform **DeepSeek-V4-Pro**, while the 1.8B model achieves better overall performance than commercial APIs such as Microsoft Translator.
+* **Real-World Scenarios and Professional Domains (WildMTBench/DomainMTBench)**: The GEMBA scores of the three models reach more than 96%–99% of Gemini 3.1 Pro (Think), and all of them outperform larger open-source models.
+* **Translation Instruction Following (IFMTBench)**: The models significantly outperform open-source models of the same scale, while the A3B model approaches the performance of Gemini 3.1 Pro (Think).
+
+In summary, Hy-MT2 is an efficient and powerful translation model series designed for complex real-world scenarios.
+
+In this release, we also open-source [IFMTBench](./IFMTBench/README.md), a benchmark for evaluating translation instruction-following capabilities.
+
+We also welcome everyone to use our released Hy-MT2-Translator Skill, which makes it easy to integrate Hy-MT2 series models for translation tasks. Download links: [ClawHub](https://clawhub.ai/tencent-adm/hy-mt2-translator-skill) and [SkillHub](https://skillhub.cn/skills/hy-mt2-translator).
+
+## News
+
+* 2026.5.21  We open-sourced **Hy-MT2-1.8B**/**Hy-MT2-7B**/**Hy-MT2-30B-A3B** on HuggingFace and ModelScope.
+* 2025.12.30 We open-sourced **HY-MT1.5-1.8B** and **HY-MT1.5-7B** on HuggingFace and ModelScope.
+* 2025.9.1 We open-sourced **Hunyuan-MT-7B** and **Hunyuan-MT-Chimera-7B** on HuggingFace and ModelScope.
+
+
+## Results
+<div align='center'>
+<img src="imgs/main_result.png" width = "100%" />
+</div>
+
+For more experimental results and analysis, please refer to our [technical report](./HY_MT2_0_Technical_Report.pdf).
+
+&nbsp;
+
+## Model Links
+| Model Name  | Description | Download Link |
+| ----------- | ----------- |-----------
+| Hy-MT2-1.8B  | Hunyuan 1.8B translation model |🤗 [Model](https://huggingface.co/tencent/Hy-MT2-1.8B)|
+| Hy-MT2-1.8B-FP8 | Hunyuan 1.8B translation model, FP8 quantization    | 🤗 [Model](https://huggingface.co/tencent/Hy-MT2-1.8B-FP8)|
+| Hy-MT2-1.8B-GGUF | Hunyuan 1.8B translation model, llama.cpp    | 🤗 [Model](https://huggingface.co/tencent/Hy-MT2-1.8B-GGUF)|
+| Hy-MT2-7B | Hunyuan 7B translation model    | 🤗 [Model](https://huggingface.co/tencent/Hy-MT2-7B)|
+| Hy-MT2-7B-FP8 | Hunyuan 7B translation model, FP8 quantization     | 🤗 [Model](https://huggingface.co/tencent/Hy-MT2-7B-FP8)|
+| Hy-MT2-7B-GGUF | Hunyuan 7B translation model, llama.cpp    | 🤗 [Model](https://huggingface.co/tencent/Hy-MT2-7B-GGUF)|
+| Hy-MT2-30B-A3B | Hunyuan 30B-A3B translation model    | 🤗 [Model](https://huggingface.co/tencent/Hy-MT2-30B-A3B)|
+| Hy-MT2-30B-A3B-FP8 | Hunyuan 30B-A3B translation model, FP8 quantization     | 🤗 [Model](https://huggingface.co/tencent/Hy-MT2-30B-A3B-FP8)|
+
+
+
+## Hy-MT2 Translation Task Instruction Examples (Chinese-English Comparison)
+
+*Note: In the following examples, both source_lang and target_lang should use the full language names. Chinese names should be used in Chinese prompts, and English names should be used in English prompts.*
+
+| Type | Chinese prompt | English prompt |
+|---|---|---|
+| **Default Translation** | 将以下文本翻译为 `{target_lang}`，注意**只需要输出翻译后的结果，不要额外解释**：<br><br>`{source_text}` | Translate the following text into `{target_lang}`. Note that you should **only output the translated result without any additional explanation**:<br><br>`{source_text}` |
+| **Terminology** | *参考下面的翻译：*<br>`{text}` 翻译成 `{text}`<br>`{text}` 翻译成 `{text}`<br>`{text}` 翻译成 `{text}`<br>将以下文本翻译为 `{target_lang}`，注意**只需要输出翻译后的结果，不要额外解释**：<br><br>`{source_text}` | *Reference the following translations:*<br>`{text}` translates to `{text}`<br>`{text}` translates to `{text}`<br>`{text}` translates to `{text}`<br><br>Translate the following text into `{target_lang}`. Note that you must **ONLY output the translated result without any additional explanation**:<br><br>`{source_text}` |
+| **Style** | 请将以下文本翻译为 `{target_lang}`。<br>注意翻译的风格要严格符合【**`{target_style}`**】<br><br>`{source_text}` | Please translate the following text into `{target_lang}`. Note that the translation style must strictly conform to [**`{target_style}`**]:<br><br>`{source_text}` |
+| **Personalization** | *【待翻译文本】*<br>`{source_text}`<br><br>*【翻译任务】*<br>1、**`{user_preferences}`**<br>2、**`{user_preferences}`**<br>3、……<br>4、将【待翻译文本】翻译为 `{target_lang}`。 | *[Source Text]*<br>`{source_text}`<br><br>*[Translation Tasks]*<br>1. **`{user_preferences}`**<br>2. **`{user_preferences}`**<br>3. ...<br>4. Translate the [Source Text] into `{target_lang}`. |
+| **Delimiters** | 请将以下文本准确翻译为 `{target_lang}`。<br>你必须在译文中**保留等量的分隔符，绝对不可遗漏、转义或翻译该符号，并注意分隔符的位置**。<br><br>`{source_text}` | Please accurately translate the following text into `{target_lang}`.<br>You must **retain the exact same number of delimiters in the translation. Strictly do not omit, escape, or translate these symbols, and pay close attention to their placement**.<br><br>`{source_text}` |
+| **Structured Data 1** | *# 任务目标*<br>将下方 `{source_text}` 中的 `{format_type}` 格式数据翻译为 `{target_lang}`。<br><br>*# 严格约束*<br>1. **结构锁定**：绝对保持原有的 `{format_type}` 数据结构、缩进和层级完全不变。<br>2. **选择性翻译**：仅翻译面向用户展示的可见文本内容。<br>3. **禁止修改**：**严禁**翻译或更改任何代码标签、键名 (Key)、变量占位符（如 `{{var}}`、`${var}`、`%s`、`%d` 等）或代码属性。<br><br>*# 数据输入*<br>`{source_text}` | *### Task*<br>Translate the user-facing text within the following `{format_type}` data into `{target_lang}`.<br><br>*### Strict Rules*<br>1. **Structure Preservation:** You MUST preserve the original `{format_type}` data structure, nesting, hierarchy, and indentation exactly as they are.<br>2. **Selective Translation:** Translate ONLY the visible, user-facing text content/values.<br>3. **Strict Non-Translation:** NEVER translate or alter code tags, keys, properties, object names, or variable placeholders. Leave them exactly in their original English/code form.<br><br>*### Source Data*<br>`{source_text}` |
+| **Structured Data 2** | *【背景信息】*<br>`{background_text}`<br><br>请结合背景信息将以下文本翻译为 `{target_lang}`。<br><br>*【待翻译文本】*<br>`{source_text}` | *[Background Information]*<br>`{background_text}`<br><br>Please translate the following text into `{target_lang}`, taking the provided background information into consideration.<br><br>*[Source Text]*<br>`{source_text}` |
+
+---
+
+## Inference and Deployment
+### transformers
+
+transformers>=5.6.0
+
+```python
+from transformers import AutoModelForCausalLM, AutoTokenizer
+import torch
+
+model_path = "tencent/Hy-MT2-30B-A3B"
+
+# Load tokenizer
+tokenizer = AutoTokenizer.from_pretrained(model_path, trust_remote_code=True)
+
+# Load model
+model = AutoModelForCausalLM.from_pretrained(
+    model_path,
+    dtype=torch.bfloat16,
+    device_map="auto",
+    trust_remote_code=True,
+)
+
+model.eval()
+
+# Example inference
+prompt = "将以下文本翻译成英语,注意只需要输出翻译后的结果,不要额外解释:\n\n今天天气真好。"
+messages = [{"role": "user", "content": prompt}]
+inputs = tokenizer.apply_chat_template(messages, add_generation_prompt=True, return_tensors="pt").to(model.device)
+
+with torch.no_grad():
+    outputs = model.generate(
+        **inputs,
+        max_new_tokens=4096,
+    )
+response = tokenizer.decode(outputs[0][inputs["input_ids"].shape[-1]:], skip_special_tokens=True)
+print(response)
+```
+
+### vllm
+
+Build vLLM from source:
+```bash
+uv venv --python 3.12 --seed --managed-python
+source .venv/bin/activate
+git clone https://github.com/vllm-project/vllm.git
+cd vllm
+uv pip install --editable . --torch-backend=auto
+```
+
+Start the vLLM server:
+
+```bash
+vllm serve tencent/Hy-MT2-30B-A3B --tensor-parallel-size 1
+```
+
+### sglang
+
+Build SGLang from source:
+```bash
+git clone https://github.com/sgl-project/sglang
+cd sglang
+pip3 install pip --upgrade
+pip3 install "transformers>=5.6.0"
+pip3 install -e "python"
+```
+
+Launch SGLang server:
+
+```bash
+python3 -m sglang.launch_server --model tencent/Hy-MT2-30B-A3B --tp 1
+```
+
+### llama_cpp
+**❕❕ This gguf depends on our STQ kernel, which is released at [PR #22836](https://github.com/ggml-org/llama.cpp/pull/22836).**
+
+#### Clone llama.cpp
+
+```bash
+git clone https://github.com/ggml-org/llama.cpp.git
+```
+
+#### Enter the llama.cpp folder
+
+```bash
+cd llama.cpp
+```
+
+#### Build llama.cpp
+
+```bash
+cmake -B build
+cmake --build build --config Release
+```
+
+#### Run a completion example
+
+```bash
+./build/bin/llama-completion \
+  --model model.gguf  \
+  -p "Translate the following segment into Chinese, without additional explanation：Hello" \
+  --jinja \
+  -ngl 0 \
+  -n 64 -st 
+```
+
+#### Run the llama.cpp benchmark
+
+```bash
+./build/bin/llama-bench -m model_zoo/model.gguf  -ngl 0
+```
+
+
+For 1.8B and 7B, we recommend using the following parameters for inference. Note that our models do not have a default system_prompt.
+
+```json
+
+{
+  "temperature": 0.7,
+  "top_p": 0.6,
+  "top_k": 20,
+  "repetition_penalty": 1.05,
+  "max_tokens": 4096
+}
+```
+
+For 30B-A3B, we recommend using the following parameters for inference. Note that our models do not have a default system_prompt.
+
+```json
+
+{
+  "temperature": 0.7,
+  "top_p": 1.0,
+  "top_k": -1,
+  "repetition_penalty": 1.0,
+  "max_tokens": 4096
+}
+```
+
+## Model Training
+Hy-MT2 provides a complete model training pipeline, supporting both full-parameter fine-tuning and LoRA fine-tuning, as well as multiple DeepSpeed ZeRO configurations and LLaMA-Factory integration.
+
+For detailed training documentation, please refer to: [Model Training Guide](./train/README.md)
+
+## Quantization Tool
+
+We provide [AngelSlim](https://github.com/tencent/AngelSlim), an easy-to-use, comprehensive, and efficient large model compression toolkit covering common quantization algorithms, low-bit quantization, speculative sampling, and more.
+
+
+## Supported Languages
+| Languages         | Abbr.   | Chinese Names   |
+|-------------------|---------|-----------------|
+| Chinese           | zh      | 中文            |
+| English           | en      | 英语            |
+| French            | fr      | 法语            |
+| Portuguese        | pt      | 葡萄牙语        |
+| Spanish           | es      | 西班牙语        |
+| Japanese          | ja      | 日语            |
+| Turkish           | tr      | 土耳其语        |
+| Russian           | ru      | 俄语            |
+| Arabic            | ar      | 阿拉伯语        |
+| Korean            | ko      | 韩语            |
+| Thai              | th      | 泰语            |
+| Italian           | it      | 意大利语        |
+| German            | de      | 德语            |
+| Vietnamese        | vi      | 越南语          |
+| Malay             | ms      | 马来语          |
+| Indonesian        | id      | 印尼语          |
+| Filipino          | tl      | 菲律宾语        |
+| Hindi             | hi      | 印地语          |
+| Traditional Chinese | zh-Hant| 繁体中文        |
+| Polish            | pl      | 波兰语          |
+| Czech             | cs      | 捷克语          |
+| Dutch             | nl      | 荷兰语          |
+| Khmer             | km      | 高棉语          |
+| Burmese           | my      | 缅甸语          |
+| Persian           | fa      | 波斯语          |
+| Gujarati          | gu      | 古吉拉特语      |
+| Urdu              | ur      | 乌尔都语        |
+| Telugu            | te      | 泰卢固语        |
+| Marathi           | mr      | 马拉地语        |
+| Hebrew            | he      | 希伯来语        |
+| Bengali           | bn      | 孟加拉语        |
+| Tamil             | ta      | 泰米尔语        |
+| Ukrainian         | uk      | 乌克兰语        |
+| Tibetan           | bo      | 藏语            |
+| Kazakh            | kk      | 哈萨克语        |
+| Mongolian         | mn      | 蒙古语          |
+| Uyghur            | ug      | 维吾尔语        |
+| Cantonese         | yue     | 粤语            |
+
+
+
+## Citing Hy-MT2
+
+```bibtex
+@misc{hy-mt1.5,
+      title={HY-MT1.5 Technical Report}, 
+      author={Mao Zheng and Zheng Li and Tao Chen and Mingyang Song and Di Wang},
+      year={2025},
+      eprint={2512.24092},
+      archivePrefix={arXiv},
+      primaryClass={cs.CL},
+      url={https://arxiv.org/abs/2512.24092}, 
+}
+```
+
+## Contact Us
+If you would like to leave feedback for our R&D and product teams, you are welcome to contact the Tencent Hunyuan LLM team. You can reach us by email at hunyuan_opensource@tencent.com.

README_CN.md

@@ -0,0 +1,297 @@
+<p align="left">
+   <a href="README.md">English</a>&nbsp;｜&nbsp;中文
+</p>
+<br>
+
+<p align="center">
+ <img src="imgs/logo-zh.png" width="400"/> <br> 
+</p>
+
+<div align="center" style="line-height: 1;">
+
+
+[![HuggingFace](https://img.shields.io/badge/%F0%9F%A4%97%20Hugging%20Face-Tencent%20Hy-ffc107?color=ffc107&logoColor=white)](https://huggingface.co/collections/tencent/hy-mt2)
+&nbsp;&nbsp;
+[![ModelScope](https://img.shields.io/badge/ModelScope-Tencent%20Hy-624aff)](https://modelscope.cn/collections/Tencent-Hunyuan/Hy-MT2)
+
+</div>
+
+<p align="center">
+    🖥️&nbsp;<a href="https://aistudio.tencent.com/"><b>官方网站</b></a>&nbsp;&nbsp;|&nbsp;&nbsp;
+    💬&nbsp;<a href="https://github.com/Tencent-Hunyuan/Hy-MT2"><b>GitHub</b></a>&nbsp;&nbsp;|&nbsp;&nbsp;
+    🪡&nbsp;<a href="https://github.com/Tencent/AngelSlim/tree/main"><b>AngelSlim</b></a></p>
+
+
+## 模型介绍
+
+
+**Hy-MT2**是涵盖Dense和MoE架构的多语言机器翻译模型系列，包含 **Hy-MT2-1.8B、7B 和 30B-A3B** 三款快思考模型，支持33种语言互译和5种民汉/方言，支持多语言指令遵循。该系列提供基于AngelSlim的**1.25-bit极致量化版**，其中1.8B模型仅需440MB存储，在苹果A15芯片上比传统4-bit推理快1.5倍。
+
+评测结果表明，Hy-MT2 在多场景下表现出众：
+
+* **通用翻译（FLORES-200）**：三款模型平均性能分别达到 **Gemini 3.1 Pro (Think)** 的 89.9%、97.9% 和 98.6%。其中 7B 和 A3B 性能超越 **DeepSeek-V4-Pro**，1.8B 综合表现超越微软翻译等商业 API。
+* **真实场景与专业领域（WildMTBench/DomainMTBench）**：三款模型 GEMBA 评分达 Gemini 3.1 Pro (Think) 的 96%~99% 以上，且均优于更大规模的开源模型。
+* **翻译指令遵循（IFMTBench）**：大幅超越同规模开源模型，A3B 性能逼近 Gemini 3.1 Pro (Think)。
+
+总之，Hy-MT2 是一个面向真实复杂场景、高效且强大的翻译模型系列。
+
+同时，本次我们也开源了一个针对翻译指令遵循能力的评测集[IFMTBench](./IFMTBench/README_zh.md)。
+
+也欢迎大家使用我们发布的 Hy-MT2-Translator Skill，可以方便接入Hy-MT2系列模型完成翻译任务，下载链接[ClawHub](https://clawhub.ai/tencent-adm/hy-mt2-translator-skill)和[SkillHub](https://skillhub.cn/skills/hy-mt2-translator)。
+
+## 新闻
+
+* 2026.5.21  我们在HuggingFace和ModelScope上开源了 **Hy-MT2-1.8B**/**Hy-MT2-7B**/**Hy-MT2-30B-A3B**/**IFMTBench**
+* 2025.12.30 我们在HuggingFace和ModelScope开源了 **HY-MT1.5-1.8B**和**HY-MT1.5-7B**
+* 2025.9.1 我们在HuggingFace和ModelScope开源了 **Hunyuan-MT-7B**和**Hunyuan-MT-Chimera-7B**。
+
+
+## 效果
+<div align='center'>
+<img src="imgs/main_result.png" width = "100%" />
+</div>
+
+更多的实验效果和分析可以参考我们的[技术报告](./HY_MT2_0_Technical_Report.pdf)。
+
+&nbsp;
+
+## 模型链接
+| 模型名  | 简介 | 下载链接 |
+| ----------- | ----------- |-----------
+| Hy-MT2-1.8B  | 混元1.8B翻译模型 |🤗 [Model](https://huggingface.co/tencent/Hy-MT2-1.8B)|
+| Hy-MT2-1.8B-FP8 | 混元1.8B翻译模型，fp8量化    | 🤗 [Model](https://huggingface.co/tencent/Hy-MT2-1.8B-FP8)|
+| Hy-MT2-1.8B-GGUF | 混元1.8B翻译模型， llama.cpp    | 🤗 [Model](https://huggingface.co/tencent/Hy-MT2-1.8B-GGUF)|
+| Hy-MT2-7B | 混元7B翻译模型    | 🤗 [Model](https://huggingface.co/tencent/Hy-MT2-7B)|
+| Hy-MT2-7B-FP8 | 混元7B翻译模型，fp8量化     | 🤗 [Model](https://huggingface.co/tencent/Hy-MT2-7B-FP8)|
+| Hy-MT2-7B-GGUF | 混元7B翻译模型， llama.cpp    | 🤗 [Model](https://huggingface.co/tencent/Hy-MT2-7B-GGUF)|
+| Hy-MT2-30B-A3B | 混元30B-A3B翻译模型    | 🤗 [Model](https://huggingface.co/tencent/Hy-MT2-30B-A3B)|
+| Hy-MT2-30B-A3B-FP8 | 混元30B-A3B翻译模型，fp8量化     | 🤗 [Model](https://huggingface.co/tencent/Hy-MT2-30B-A3B-FP8)|
+
+
+
+## Hy-MT2 翻译任务指令示例（中英文对照）
+
+*注意：下面的source_lang和target_lang都使用语言的全称，中文使用中文全称，英文使用英文全称。*
+
+| Type | Chinese prompt | English prompt |
+|---|---|---|
+| **Default Translation** | 将以下文本翻译为 `{target_lang}`，注意**只需要输出翻译后的结果，不要额外解释**：<br><br>`{source_text}` | Translate the following text into `{target_lang}`. Note that you should **only output the translated result without any additional explanation**:<br><br>`{source_text}` |
+| **Terminology** | *参考下面的翻译：*<br>`{text}` 翻译成 `{text}`<br>`{text}` 翻译成 `{text}`<br>`{text}` 翻译成 `{text}`<br>将以下文本翻译为 `{target_lang}`，注意**只需要输出翻译后的结果，不要额外解释**：<br><br>`{source_text}` | *Reference the following translations:*<br>`{text}` translates to `{text}`<br>`{text}` translates to `{text}`<br>`{text}` translates to `{text}`<br><br>Translate the following text into `{target_lang}`. Note that you must **ONLY output the translated result without any additional explanation**:<br><br>`{source_text}` |
+| **Style** | 请将以下文本翻译为 `{target_lang}`。<br>注意翻译的风格要严格符合【**`{target_style}`**】<br><br>`{source_text}` | Please translate the following text into `{target_lang}`. Note that the translation style must strictly conform to [**`{target_style}`**]:<br><br>`{source_text}` |
+| **Personalization** | *【待翻译文本】*<br>`{source_text}`<br><br>*【翻译任务】*<br>1、**`{user_preferences}`**<br>2、**`{user_preferences}`**<br>3、……<br>4、将【待翻译文本】翻译为 `{target_lang}`。 | *[Source Text]*<br>`{source_text}`<br><br>*[Translation Tasks]*<br>1. **`{user_preferences}`**<br>2. **`{user_preferences}`**<br>3. ...<br>4. Translate the [Source Text] into `{target_lang}`. |
+| **Delimiters** | 请将以下文本准确翻译为 `{target_lang}`。<br>你必须在译文中**保留等量的分隔符，绝对不可遗漏、转义或翻译该符号，并注意分隔符的位置**。<br><br>`{source_text}` | Please accurately translate the following text into `{target_lang}`.<br>You must **retain the exact same number of delimiters in the translation. Strictly do not omit, escape, or translate these symbols, and pay close attention to their placement**.<br><br>`{source_text}` |
+| **Structured Data 1** | *# 任务目标*<br>将下方 `{source_text}` 中的 `{format_type}` 格式数据翻译为 `{target_lang}`。<br><br>*# 严格约束*<br>1. **结构锁定**：绝对保持原有的 `{format_type}` 数据结构、缩进和层级完全不变。<br>2. **选择性翻译**：仅翻译面向用户展示的可见文本内容。<br>3. **禁止修改**：**严禁**翻译或更改任何代码标签、键名 (Key)、变量占位符（如 `{{var}}`、`${var}`、`%s`、`%d` 等）或代码属性。<br><br>*# 数据输入*<br>`{source_text}` | *### Task*<br>Translate the user-facing text within the following `{format_type}` data into `{target_lang}`.<br><br>*### Strict Rules*<br>1. **Structure Preservation:** You MUST preserve the original `{format_type}` data structure, nesting, hierarchy, and indentation exactly as they are.<br>2. **Selective Translation:** Translate ONLY the visible, user-facing text content/values.<br>3. **Strict Non-Translation:** NEVER translate or alter code tags, keys, properties, object names, or variable placeholders. Leave them exactly in their original English/code form.<br><br>*### Source Data*<br>`{source_text}` |
+| **Structured Data 2** | *【背景信息】*<br>`{background_text}`<br><br>请结合背景信息将以下文本翻译为 `{target_lang}`。<br><br>*【待翻译文本】*<br>`{source_text}` | *[Background Information]*<br>`{background_text}`<br><br>Please translate the following text into `{target_lang}`, taking the provided background information into consideration.<br><br>*[Source Text]*<br>`{source_text}` |
+
+---
+
+## 推理和部署
+### transformers
+
+transformers>=5.6.0
+
+```python
+from transformers import AutoModelForCausalLM, AutoTokenizer
+import torch
+
+model_path = "tencent/Hy-MT2-30B-A3B"
+
+# Load tokenizer
+tokenizer = AutoTokenizer.from_pretrained(model_path, trust_remote_code=True)
+
+# Load model
+model = AutoModelForCausalLM.from_pretrained(
+    model_path,
+    dtype=torch.bfloat16,
+    device_map="auto",
+    trust_remote_code=True,
+)
+
+model.eval()
+
+# Example inference
+prompt = "将以下文本翻译成英语,注意只需要输出翻译后的结果,不要额外解释:\n\n今天天气真好。"
+messages = [{"role": "user", "content": prompt}]
+inputs = tokenizer.apply_chat_template(messages, add_generation_prompt=True, return_tensors="pt").to(model.device)
+
+with torch.no_grad():
+    outputs = model.generate(
+        **inputs,
+        max_new_tokens=4096,
+    )
+response = tokenizer.decode(outputs[0][inputs["input_ids"].shape[-1]:], skip_special_tokens=True)
+print(response)
+```
+
+### vllm
+
+Build vLLM from source:
+```bash
+uv venv --python 3.12 --seed --managed-python
+source .venv/bin/activate
+git clone https://github.com/vllm-project/vllm.git
+cd vllm
+uv pip install --editable . --torch-backend=auto
+```
+
+Start the vLLM server:
+
+```bash
+vllm serve tencent/Hy-MT2-30B-A3B --tensor-parallel-size 1
+```
+
+### sglang
+
+Build SGLang from source:
+```bash
+git clone https://github.com/sgl-project/sglang
+cd sglang
+pip3 install pip --upgrade
+pip3 install "transformers>=5.6.0"
+pip3 install -e "python"
+```
+
+Launch SGLang server:
+
+```bash
+python3 -m sglang.launch_server --model tencent/Hy-MT2-30B-A3B --tp 1
+```
+
+
+### llama_cpp
+**❕❕ This gguf depends on our STQ kernel, which is released at [PR #22836](https://github.com/ggml-org/llama.cpp/pull/22836).**
+
+#### Clone llama.cpp
+
+```bash
+git clone https://github.com/ggml-org/llama.cpp.git
+```
+
+#### Enter the llama.cpp folder
+
+```bash
+cd llama.cpp
+```
+
+#### Build llama.cpp
+
+```bash
+cmake -B build
+cmake --build build --config Release
+```
+
+#### Run a completion example
+
+```bash
+./build/bin/llama-completion \
+  --model model.gguf  \
+  -p "Translate the following segment into Chinese, without additional explanation：Hello" \
+  --jinja \
+  -ngl 0 \
+  -n 64 -st 
+```
+
+#### Run the llama.cpp benchmark
+
+```bash
+./build/bin/llama-bench -m model_zoo/model.gguf  -ngl 0
+```
+
+
+对于1.8B和7B，我们推荐使用下面这组参数进行推理。注意，我们的模型没有默认 system_prompt。
+
+```json
+
+{
+  "temperature": 0.7,
+  "top_p": 0.6,
+  "top_k": 20,
+  "repetition_penalty": 1.05,
+  "max_tokens": 4096
+}
+```
+
+对于30B-A3B，我们推荐使用下面这组参数进行推理。注意，我们的模型没有默认 system_prompt。
+
+```json
+
+{
+  "temperature": 0.7,
+  "top_p": 1.0,
+  "top_k": -1,
+  "repetition_penalty": 1.0,
+  "max_tokens": 4096
+}
+```
+
+
+## 模型训练
+Hy-MT2提供了完整的模型训练流程，支持全量微调和 LoRA 微调，同时支持 DeepSpeed ZeRO 多种配置以及 LLaMA-Factory 集成。
+
+详细的训练文档请参考：[模型训练指南](./train/README_CN.md)
+
+## 量化工具
+
+我们提供了 [AngelSlim](https://github.com/tencent/AngelSlim)——一套易用、全面、高效的大模型压缩工具包，涵盖常用量化算法、低比特量化和投机采样等能力。
+
+
+## 支持的语种
+| Languages         | Abbr.   | Chinese Names   |
+|-------------------|---------|-----------------|
+| Chinese           | zh      | 中文            |
+| English           | en      | 英语            |
+| French            | fr      | 法语            |
+| Portuguese        | pt      | 葡萄牙语        |
+| Spanish           | es      | 西班牙语        |
+| Japanese          | ja      | 日语            |
+| Turkish           | tr      | 土耳其语        |
+| Russian           | ru      | 俄语            |
+| Arabic            | ar      | 阿拉伯语        |
+| Korean            | ko      | 韩语            |
+| Thai              | th      | 泰语            |
+| Italian           | it      | 意大利语        |
+| German            | de      | 德语            |
+| Vietnamese        | vi      | 越南语          |
+| Malay             | ms      | 马来语          |
+| Indonesian        | id      | 印尼语          |
+| Filipino          | tl      | 菲律宾语        |
+| Hindi             | hi      | 印地语          |
+| Traditional Chinese | zh-Hant| 繁体中文        |
+| Polish            | pl      | 波兰语          |
+| Czech             | cs      | 捷克语          |
+| Dutch             | nl      | 荷兰语          |
+| Khmer             | km      | 高棉语          |
+| Burmese           | my      | 缅甸语          |
+| Persian           | fa      | 波斯语          |
+| Gujarati          | gu      | 古吉拉特语      |
+| Urdu              | ur      | 乌尔都语        |
+| Telugu            | te      | 泰卢固语        |
+| Marathi           | mr      | 马拉地语        |
+| Hebrew            | he      | 希伯来语        |
+| Bengali           | bn      | 孟加拉语        |
+| Tamil             | ta      | 泰米尔语        |
+| Ukrainian         | uk      | 乌克兰语        |
+| Tibetan           | bo      | 藏语            |
+| Kazakh            | kk      | 哈萨克语        |
+| Mongolian         | mn      | 蒙古语          |
+| Uyghur            | ug      | 维吾尔语        |
+| Cantonese         | yue     | 粤语            |
+
+
+
+## Citing Hy-MT2
+
+```bibtex
+@misc{hy-mt1.5,
+      title={HY-MT1.5 Technical Report}, 
+      author={Mao Zheng and Zheng Li and Tao Chen and Mingyang Song and Di Wang},
+      year={2025},
+      eprint={2512.24092},
+      archivePrefix={arXiv},
+      primaryClass={cs.CL},
+      url={https://arxiv.org/abs/2512.24092}, 
+}
+```
+
+## 联系我们
+如果你想给我们的研发和产品团队留言，欢迎联系我们腾讯混元LLM团队。你可以通过邮件（hunyuan_opensource@tencent.com）联系我们。

train/README.md

@@ -0,0 +1,303 @@
+<p align="left">
+    <a href="README_CN.md">中文</a> ｜ English
+</p>
+
+# Model Training
+
+Hy3 preview provides processes related to model training. This section details how to process training data for model training purposes.
+
+## Training Data Format and Processing
+
+The training data should be formatted as a list of messages. By default, the system prompt for both training and inference is empty, but you may customize it as needed.
+
+Below is a training data example for a translation task:
+
+```python
+# Translation task example
+{"messages": [{"role": "user", "content": "将以下中文翻译为英文，只输出翻译结果，不要额外解释：\n\n实验结果证明了假设的正确性。"}, {"role": "assistant", "content": "The experimental results demonstrate the correctness of the hypothesis."}]}
+
+```
+
+## Quick Start
+
+You can quickly get started by following the instructions in the Quick Start Guide.
+
+## Model Training
+
+### Hardware Requirements
+
+The following are the minimum hardware requirements for each model at max_seq_length = 8192:
+
+#### Hy-MT2-1.8B (Dense)
+
+| Training Method | DeepSpeed Strategy | Minimum GPU Requirement |
+|----------------|-------------------|------------------------|
+| LoRA Fine-tuning | ZeRO-2 (no offload) | 1 GPU (24GB+) |
+| Full Fine-tuning | ZeRO-2 (no offload) | 1 GPU (24GB+) |
+
+#### Hy-MT2-7B (Dense)
+
+| Training Method | DeepSpeed Strategy | Minimum GPU Requirement |
+|----------------|-------------------|------------------------|
+| LoRA Fine-tuning | ZeRO-2 (no offload) | 1 GPU (80GB+) |
+| Full Fine-tuning | ZeRO-3 (no offload) | 2 GPUs (80GB+ each) |
+
+#### Hy-MT2-30B-A3B (MoE)
+
+| Training Method | DeepSpeed Strategy | Minimum GPU Requirement |
+|----------------|-------------------|------------------------|
+| LoRA Fine-tuning | ZeRO-2 (no offload) | 8 GPUs on a single machine (80GB+ each) |
+| Full Fine-tuning | ZeRO-3 + offload | 8 GPUs on a single machine (80GB+ each) |
+
+### Configure Passwordless SSH Login Between Machines (Multi-Machine Training)
+
+> If you only use single-machine training, you can skip this section.
+
+The following instructions use two machines as an example, with their IPs denoted as `${ip1}` and `${ip2}`. All steps should be performed inside the Docker container.
+
+First, configure passwordless SSH for each container on every machine:
+
+```sh
+ssh-keygen			# Generate id_rsa and id_rsa.pub for passwordless login
+ssh-keygen -t rsa -A    # Generate /etc/ssh/ssh_host_rsa_key and ssh_host_ecdsa_key for SSH listening
+/usr/sbin/sshd -p 36005 -o ListenAddress=0.0.0.0        # Start SSH listening
+echo "Port 36005" > ~/.ssh/config   # Set SSH connection port to 36005
+passwd root    # Set the root password to avoid monitoring platform alerts
+```
+
+Note: `36005` is an example port. You may use any available port, but ensure it is **open** and **not occupied by other processes**.
+
+Next, in each machine's container, execute:
+
+```sh
+cat ~/.ssh/id_rsa.pub
+```
+
+**Copy the output SSH public key and paste it into the `~/.ssh/authorized_keys` file, one key per line. This must be done on every machine.** In the end, the `~/.ssh/authorized_keys` file on each machine should be identical and contain the public keys of all machines.
+
+Please note that for multi-node training, the code executed on each node must be identical. It is recommended to mount a shared network drive. If this is not possible, you must manually copy the dataset, scripts, and code to the same directory on each machine.
+
+### Launch Methods
+
+This project provides two training methods. You can choose based on your needs:
+
+- **DeepSpeed Native Training** (based on HuggingFace Transformers Trainer): Located in the `train/deepspeed_support` directory
+- **LLaMA-Factory Training**: Located in the `train/llama_factory_support` directory
+
+#### DeepSpeed Native Training
+
+Reference: [HuggingFace Transformers Trainer](https://huggingface.co/docs/transformers/main/en/main_classes/trainer)
+
+##### Training Scripts
+
+In the `train/deepspeed_support` directory, the training scripts for each model are as follows:
+
+| Model | Full Fine-tuning | LoRA Fine-tuning |
+|-------|-----------------|-----------------|
+| Hy-MT2-1.8B (Dense) | `bash train_dense.sh 1.8B` | `bash train_dense_lora.sh 1.8B` |
+| Hy-MT2-7B (Dense) | `bash train_dense.sh 7B` | `bash train_dense_lora.sh 7B` |
+| Hy-MT2-30B-A3B (MoE) | `bash train.sh` | `bash train_lora.sh` |
+
+##### Single-Machine Training
+
+In the `train/deepspeed_support` directory, install dependencies and execute the corresponding script:
+
+```sh
+pip install -r requirements.txt
+# Example: Dense 1.8B full fine-tuning
+bash train_dense.sh 1.8B
+```
+
+##### Multi-Machine Training
+
+To launch training across multiple machines, please first complete the configuration in [Configure Passwordless SSH Login Between Machines](#configure-passwordless-ssh-login-between-machines-multi-machine-training), and ensure all machines are within the same cluster.
+
+Confirm that dependencies are installed (if not, run `pip install -r requirements.txt`), then set the `IP_LIST` environment variable in the corresponding training script:
+
+```shell
+export HOST_GPU_NUM=8
+# IP list, comma separated. e.g. "192.168.1.1,192.168.1.2" or single node "192.168.1.1"
+IP_LIST=${IP_LIST:-"127.0.0.1"}
+```
+
+Note: If the `IP_LIST` environment variable is not set, replace `IP_LIST` with the IP list! The format is:
+```
+For a single IP:
+IP_LIST=${ip_1}
+
+For multiple IPs:
+IP_LIST=${ip_1},${ip_2}
+
+```
+
+Replace `${ip_1}` and `${ip_2}` with the actual IP addresses.
+
+Then, on the machine with `${ip1}`, execute the corresponding training script in the `train/deepspeed_support/` directory. On first launch, you may see the following output:
+
+```ssh
+The authenticity of host '[ip]:36005 ([ip]:36005)' can't be established.
+ECDSA key fingerprint is xxxxxx.
+ECDSA key fingerprint is MD5:xxxxxx.
+Are you sure you want to continue connecting (yes/no)?
+```
+
+Type `yes` to continue.
+
+##### Key Parameters
+
+The key parameters in the script are as follows:
+
+- `--deepspeed`: Path to the DeepSpeed configuration file. Three default DeepSpeed configuration files are provided in the `train/deepspeed_support` folder: `ds_zero2_no_offload.json`, `ds_zero3_no_offload.json`, and `ds_zero3_offload.json`, with decreasing memory requirements in that order.
+- `--model_name_or_path`: Path to the Hy3 preview HF pre-trained model weights to load.
+- `--tokenizer_name_or_path`: Path to the tokenizer folder.
+- `--train_data_file`: Path to the training file, which should be a jsonl file.
+- `--output_dir`: Output directory where logs, tensorboard files, and model weights will be stored.
+- `--per_device_train_batch_size`: Batch size per GPU.
+- `--gradient_accumulation_steps`: Number of gradient accumulation steps. The global batch size is `per_device_train_batch_size * gradient_accumulation_steps * dp_size`.
+- `--max_steps`: Total number of training steps.
+- `--save_steps`: Number of steps between saving checkpoints.
+- `--use_lora`: Whether to use LoRA training. Also accepts `--lora_rank`, `--lora_alpha`, and `--lora_dropout` parameters. By default, LoRA is applied to "q_proj", "k_proj", "v_proj", and "o_proj". To change this, modify the code. Note: **When using LoRA training, only the LoRA weights are saved, not the base model weights.** To merge LoRA weights, see the "LoRA Weight Merging" section below.
+- `--make_moe_param_leaf_module`: When using ZeRO-3 with MoE training, treat the MoE module as a leaf module, i.e., its parameters are not partitioned by ZeRO-3. This option is expected to significantly increase memory usage.
+- `--gradient_checkpointing`: Enable gradient checkpointing.
+- `--train_attention_params_only`: Whether to train only attention parameters.
+- `--learning_rate`: Maximum learning rate during training.
+- `--min_lr`: Minimum learning rate during training.
+- `--use_flash_attn`: Enable flash-attention for accelerated training.
+
+**Notes:**
+
+- To resume training from a previously saved checkpoint rather than loading pre-trained weights, specify `--resume_from_checkpoint` with the path to the checkpoint. Do not specify `--model_name_or_path`; this will load only the weights without the training state.
+- When resuming from a checkpoint, there may be minor differences in loss due to the randomness of some non-deterministic algorithms. This is normal. See: [HuggingFace Transformers Trainer Randomness](https://huggingface.co/docs/transformers/main/en/main_classes/trainer#randomness)
+- When `--model_name_or_path` is specified, all model-related parameters will be ignored.
+- Samples within a batch are padded to the length of the longest sample in the batch, but the maximum length of each sample is `max_seq_length`. Any excess will be truncated.
+- If you see a warning about bias weights not being loaded, you can ignore it. Hunyuan-Large does not use bias.
+
+##### What if GPU Memory is Insufficient?
+
+Reference: [DeepSpeed Configuration](https://www.deepspeed.ai/docs/config-json/)
+
+You can try modifying the DeepSpeed configuration by removing the `auto` attribute from the following parameters and reducing their values:
+
+- `stage3_param_persistence_threshold`
+- `stage3_prefetch_bucket_size`
+- `stage3_max_reuse_distance`
+
+##### LoRA Weight Merging
+
+LoRA weights saved during training cannot be merged into the ZeRO-3 model at runtime, as ZeRO-3 partitions model weights across data parallel ranks. To merge LoRA weights into the base model, you can do so offline to obtain a merged weight file. Run `merge_lora_weight.sh` to merge the LoRA and base model weights. The parameters are:
+
+- `--base_model_path`: Directory of the base model weights
+- `--adapter_model_path`: Directory of the LoRA weights
+- `--output_path`: Directory to save the merged weights
+- `--save_dtype`: Data type for saving the merged weights; options are: fp16, bf16, fp32
+
+#### LLaMA-Factory Training
+
+If you are familiar with LLaMA-Factory, you may use it for fine-tuning. All scripts, code, and configuration files are archived in the `train/llama_factory_support` directory. Unless otherwise specified, all files mentioned below are located in this directory.
+
+##### Installation
+
+You can install LLaMA-Factory by downloading the source code from https://github.com/hiyouga/LLaMA-Factory/tree/main and following the instructions on the website.
+
+##### Training Scripts and Configuration Files
+
+The configuration files and launch scripts for each model are as follows:
+
+| Model | Full Fine-tuning Config | LoRA Fine-tuning Config | Launch Script |
+|-------|------------------------|------------------------|---------------|
+| Hy-MT2-1.8B (Dense) | `hy_dense_1_8b_full_sft.yaml` | `hy_dense_1_8b_lora_sft.yaml` | `bash train_lf_dense.sh` |
+| Hy-MT2-7B (Dense) | `hy_dense_7b_full_sft.yaml` | `hy_dense_7b_lora_sft.yaml` | `YAML_FILE=hy_dense_7b_full_sft.yaml bash train_lf_dense.sh` |
+| Hy-MT2-30B-A3B (MoE) | `hy_v3_full_sft.yaml` | `hy_v3_lora_sft.yaml` | `bash train_lf.sh` |
+
+> **Tip**: The Dense model launch script `train_lf_dense.sh` uses `hy_dense_1_8b_full_sft.yaml` by default. You can specify other configuration files via the `YAML_FILE` environment variable.
+
+Key parameters in the configuration files are as follows:
+
+**Model:**
+
+- `model_name_or_path`: Path to the Hy-MT HF format pre-trained model weights
+- `trust_remote_code`: Whether to trust remote code; Hy-MT requires this to be set to `true`
+
+**Training Method:**
+
+- `stage`: Training stage, currently `sft` (supervised fine-tuning)
+- `finetuning_type`: Fine-tuning type, either `full` (full fine-tuning) or `lora` (LoRA fine-tuning)
+- `deepspeed`: DeepSpeed configuration file path; `ds_zero3_offload.json` is recommended for full fine-tuning, `ds_zero2_offload_lora.json` for LoRA fine-tuning
+
+**LoRA Parameters (only effective during LoRA fine-tuning):**
+
+- `lora_rank`: LoRA rank, default `64`
+- `lora_alpha`: LoRA alpha coefficient, default `128`
+- `lora_dropout`: LoRA dropout ratio, default `0.05`
+- `lora_target`: Target modules for LoRA, default `q_proj,k_proj,v_proj,o_proj`
+
+**Dataset:**
+
+- `dataset_dir`: Dataset directory path
+- `dataset`: Dataset name, must be registered in `dataset_info.json` under `dataset_dir`
+- `template`: Chat template; Hy-MT2-1.8B uses `hy_dense_1_8b`, Hy-MT2-7B uses `hy_dense_7b`, Hy-MT2-30B-A3B uses `hy_v3`
+- `cutoff_len`: Maximum sequence length; sequences exceeding this will be truncated. For full fine-tuning, can be set to `262144` (262K); for LoRA fine-tuning, `8192` is recommended to save memory
+- `max_samples`: Maximum number of samples per dataset
+- `overwrite_cache`: Whether to overwrite cached preprocessed datasets
+
+**Output:**
+
+- `output_dir`: Output directory where logs, TensorBoard files, and weights will be stored
+- `logging_steps`: Number of steps between logging
+- `save_steps`: Number of steps between saving checkpoints
+- `plot_loss`: Whether to plot the training loss curve
+- `overwrite_output_dir`: Whether to overwrite the existing output directory
+- `save_only_model`: Whether to save only model weights (excluding optimizer states, etc.)
+- `report_to`: Logging tool, options: `none`, `wandb`, `tensorboard`, `swanlab`, `mlflow`
+
+**Training Hyperparameters:**
+
+- `per_device_train_batch_size`: Batch size per GPU
+- `gradient_accumulation_steps`: Gradient accumulation steps; `per_device_train_batch_size * gradient_accumulation_steps * dp_size` equals the global batch size
+- `learning_rate`: Maximum learning rate; `1.0e-5` recommended for full fine-tuning, `2.0e-4` for LoRA fine-tuning
+- `num_train_epochs`: Number of training epochs
+- `lr_scheduler_type`: Learning rate scheduler type; `cosine_with_min_lr` is recommended
+- `lr_scheduler_kwargs.min_lr_rate`: Ratio of minimum to maximum learning rate; e.g., `0.1` means the minimum learning rate is 10% of the maximum
+- `warmup_ratio`: Proportion of total training steps used for warmup
+- `bf16`: Whether to use BFloat16 mixed precision training
+- `gradient_checkpointing`: Whether to enable gradient checkpointing to save memory
+- `ddp_timeout`: Distributed training timeout (milliseconds)
+- `flash_attn`: Attention implementation; `fa2` (FlashAttention-2) is recommended, `sdpa` is also available; using `fa2` requires the flash-attn package
+- `resume_from_checkpoint`: Resume training from a specified checkpoint path; set to `null` to start from scratch
+
+##### Launch Training
+
+For multi-machine training, please first complete the configuration in [Configure Passwordless SSH Login Between Machines](#configure-passwordless-ssh-login-between-machines-multi-machine-training) (single-machine training can skip this step).
+
+Modify the following configuration at the beginning of the corresponding launch script:
+
+```shell
+export HOST_GPU_NUM=8
+# IP list, comma separated. e.g. "192.168.1.1,192.168.1.2" or single node "192.168.1.1"
+export IP_LIST=${IP_LIST:-"127.0.0.1"}
+```
+
+Note: If the `IP_LIST` environment variable is not set, replace `IP_LIST` with the IP list! The format is:
+```
+For a single IP:
+IP_LIST=${ip_1}
+
+For multiple IPs:
+IP_LIST=${ip_1},${ip_2}
+
+```
+
+Replace `${ip_1}` and `${ip_2}` with the actual IP addresses.
+
+Then, on each machine, run the corresponding launch script in the `train/llama_factory_support/` directory. For example:
+
+```sh
+# Dense 1.8B full fine-tuning
+bash train_lf_dense.sh
+
+# Dense 7B LoRA fine-tuning
+YAML_FILE=hy_dense_7b_lora_sft.yaml bash train_lf_dense.sh
+
+# MoE 30B-A3B full fine-tuning
+bash train_lf.sh
+```

train/README_CN.md

@@ -0,0 +1,303 @@
+<p align="left">
+   <a href="README.md">English</a> ｜ 中文
+</p>
+
+# 模型训练
+
+Hy-MT 提供了模型训练相关流程，您可以在此章节对训练数据格式进行处理以供模型训练使用。
+
+## 训练数据格式及处理
+
+训练数据按照以下形式处理为 messages 格式，训练和推理的默认 system prompt 为空，可以根据自己的需求进行设定。
+
+以下是翻译任务的训练数据示例：
+
+```python
+# 翻译任务示例
+{"messages": [{"role": "user", "content": "将以下中文翻译为英文，只输出翻译结果，不要额外解释：\n\n实验结果证明了假设的正确性。"}, {"role": "assistant", "content": "The experimental results demonstrate the correctness of the hypothesis."}]}
+
+```
+
+## 快速开始
+
+您可以参照快速开始文档中的内容进行快速上手。
+
+## 模型训练
+
+### 硬件需求
+
+以下为各模型在 max_seq_length = 8192 时的最低硬件需求：
+
+#### Hy-MT2-1.8B（Dense）
+
+| 训练方式 | DeepSpeed 策略 | 最低 GPU 需求 |
+|---------|---------------|-------------|
+| LoRA 微调 | ZeRO-2（无 offload） | 1 卡（24GB+） |
+| 全量微调 | ZeRO-2（无 offload） | 1 卡（24GB+） |
+
+#### Hy-MT2-7B（Dense）
+
+| 训练方式 | DeepSpeed 策略 | 最低 GPU 需求 |
+|---------|---------------|-------------|
+| LoRA 微调 | ZeRO-2（无 offload） | 1 卡（80GB+） |
+| 全量微调 | ZeRO-3（无 offload） | 2 卡（80GB+ 每卡） |
+
+#### Hy-MT2-30B-A3B（MoE）
+
+| 训练方式 | DeepSpeed 策略 | 最低 GPU 需求 |
+|---------|---------------|-------------|
+| LoRA 微调 | ZeRO-2（无 offload） | 单机 8 卡（80GB+ 每卡） |
+| 全量微调 | ZeRO-3 + offload | 单机 8 卡（80GB+ 每卡） |
+
+### 配置机器间免密 ssh 登录（多机训练）
+
+> 如果只使用单机训练，可跳过本节。
+
+以下操作以两个机器为例，两台机器的 ip 分别以`${ip1}`和`${ip2}`标识，以下操作均在 docker container 内执行。
+
+首先，配置多机container免密，在每台机器上执行。
+
+```sh
+ssh-keygen			# 生成id_rsa和id_rsa.pub，用于免密登录
+ssh-keygen -t rsa -A    # 生成/etc/ssh/ssh_host_rsa_key和ssh_host_ecdsa_key， 用于后面启动ssh listen
+/usr/sbin/sshd -p 36005 -o ListenAddress=0.0.0.0        # 启动 SSH 监听
+echo "Port 36005" > ~/.ssh/config   # ssh 连接端口修改为 36005
+passwd root    # 需要配置root密码，否则监测平台会报警
+```
+
+注意：这里的`36005`是一个示例端口，可以选用任意端口，但需要保证使用的端口**开放**且**不被其他的进程占用**。
+
+接下来，在每台机器的 container 内，执行：
+
+```sh
+cat ~/.ssh/id_rsa.pub
+```
+
+**将输出的 ssh 公钥复制并粘贴到`~/.ssh/authorized_keys`文件中，每行一个公钥，每台机器上都要做这个操作**。最终每台机器上的`~/.ssh/authorized_keys`文件内容应当是一致的，并且包含了所有机器的公钥。
+
+需要注意，多节点训练时，每个节点上执行的代码都得一致，建议挂载一个共享的网络盘，如果无法挂载共享网盘，则需要手动将数据集、脚本、代码复制在多台机器的相同目录下。
+
+### 启动方式
+
+本项目提供两种训练方式，您可以根据需求选择：
+
+- **DeepSpeed 原生训练**（基于 HuggingFace Transformers Trainer）：位于 `train/deepspeed_support` 目录下
+- **LLaMA-Factory 训练**：位于 `train/llama_factory_support` 目录下
+
+#### DeepSpeed 原生训练
+
+参考：[HuggingFace Transformers Trainer](https://huggingface.co/docs/transformers/main/en/main_classes/trainer)
+
+##### 训练脚本
+
+在 `train/deepspeed_support` 目录下，各模型对应的训练脚本如下：
+
+| 模型 | 全量微调 | LoRA 微调 |
+|------|---------|----------|
+| Hy-MT2-1.8B（Dense） | `bash train_dense.sh 1.8B` | `bash train_dense_lora.sh 1.8B` |
+| Hy-MT2-7B（Dense） | `bash train_dense.sh 7B` | `bash train_dense_lora.sh 7B` |
+| Hy-MT2-30B-A3B（MoE） | `bash train.sh` | `bash train_lora.sh` |
+
+##### 单机启动训练
+
+在 `train/deepspeed_support` 目录下，安装依赖后直接执行对应的脚本命令即可：
+
+```sh
+pip install -r requirements.txt
+# 示例：Dense 1.8B 全量微调
+bash train_dense.sh 1.8B
+```
+
+##### 多机启动训练
+
+如果要用多台机器启动训练，请先完成 [配置机器间免密 ssh 登录](#配置机器间免密-ssh-登录多机训练) 中的配置，并保证多台机器在一个集群内。
+
+确认依赖已经安装完成（如未安装，请执行`pip install -r requirements.txt`安装），然后在对应的训练脚本中设置 `IP_LIST` 环境变量：
+
+```shell
+export HOST_GPU_NUM=8
+# IP list, comma separated. e.g. "192.168.1.1,192.168.1.2" or single node "192.168.1.1"
+IP_LIST=${IP_LIST:-"127.0.0.1"}
+```
+
+注意：如果`IP_LIST`环境变量未设置，则将`IP_LIST`替换为IP列表！格式为：
+```
+如果只有一个IP：
+IP_LIST=${ip_1}
+
+如果有多个IP：
+IP_LIST=${ip_1},${ip_2}
+
+```
+
+请将`${ip_1}`和`${ip_2}`替换为真实的IP地址。
+
+然后，在`${ip1}`的机器上，在`train/deepspeed_support/`目录下，执行对应的训练脚本即可。注意第一次启动时可能会看见以下的输出：
+
+```ssh
+The authenticity of host '[ip]:36005 ([ip]:36005)' can't be established.
+ECDSA key fingerprint is xxxxxx.
+ECDSA key fingerprint is MD5:xxxxxx.
+Are you sure you want to continue connecting (yes/no)?
+```
+
+此时输入`yes`即可继续。
+
+##### 关键参数
+
+脚本中的关键参数如下：
+
+- `--deepspeed`: 此参数应当指向一个 deepspeed 的配置文件，`train/deepspeed_support`文件夹下提供了三种 DeepSpeed 的默认配置文件：`ds_zero2_no_offload.json`, `ds_zero3_no_offload.json`, `ds_zero3_offload.json`，这三个配置文件所需显存依次减少
+- `--model_name_or_path`: 要加载的 Hy-MT 的 HF 预训练模型权重，否则无法加载
+- `--tokenizer_name_or_path`: tokenizer 文件夹路径, 否则无法加载
+- `--train_data_file`: 训练文件路径，应该为一个 jsonl 文件
+- `--output_dir`: 输出文件夹，log、tensorboard 和权重都会存储在这个路径下
+- `--per_device_train_batch_size`: 每张卡上的 batch size
+- `--gradient_accumulation_steps`: 梯度累计次数，`per_device_train_batch_size * gradient_accumulation_steps * dp_size`为 global_batch_size
+- `--max_steps`: 训练的总步数
+- `--save_steps`: 每多少个 step 存储一个 checkpoint
+- `--use_lora`: 是否用 lora 训练，同时接收`--lora_rank`，`--lora_alpha`和`--lora_dropout`参数。lora 默认应用于 "q_proj", "k_proj", "v_proj", "o_proj" 四个参数，如果需要改变的话在代码中修改即可。注意：**使用 lora 训练时，只会保存 lora 的权重，而不会保存 base 模型的权重**，如果需要合并 lora 权重，看下面的“Lora 权重合并”一节
+- `--make_moe_param_leaf_module`：当用 zero3 以及 MoE 训练时，将 MoE 模块视作一个 leaf module，即它的参数不进行 zero3 切分，这个选项预计会显著增加显存占用
+- `--gradient_checkpointing`：开启梯度重计算
+- `--train_attention_params_only`: 是否只训练 attention 参数
+- `--learning_rate`: 训练时的最大学习率
+- `--min_lr`: 训练时的最小学习率
+- `--use_flash_attn`: 开启 flash-attention 进行训练加速
+
+**注意：**
+
+- 如果想从一个中途保存的 ckpt 继续训练，而不是加载一个预训练的权重，直接指定`--resume_from_checkpoint`为之前训练保存的 ckpt 路径，不要指定`--model_name_or_path`，这样只会加载权重，而不会加载训练状态
+- 从 ckpt 继续训练时，loss 可能会有微小的偏差，这是由一些非确定性算法带来的随机性，是正常现象。参考：[HuggingFace Transformers Trainer Randomness](https://huggingface.co/docs/transformers/main/en/main_classes/trainer#randomness)
+- 当 `--model_name_or_path` 有效时，所有模型相关的参数都会被忽略
+- 一个 batch 内的样本会通过 padding 对齐 batch 内最长的样本，而每条样本的长度最长为 max_seq_length，超出的部分会被裁剪
+- 如果报出 bias 权重没有 load 的 warning，忽略即可，Hunyuan-Large 中不会用到 bias
+
+##### 显存不足怎么办？
+
+参考：[DeepSpeed Configuration](https://www.deepspeed.ai/docs/config-json/)
+
+可以尝试修改 ds config，去掉这几个参数的 auto 属性，改小试试看：
+
+- `stage3_param_persistence_threshold`
+- `stage3_prefetch_bucket_size`
+- `stage3_max_reuse_distance`
+
+##### Lora 模型合并
+
+保存下来的 lora 权重没法在训练运行时合并到 zero3 模型中，因为 zero3 开启时模型权重会切分到各 dp rank 上。因此如果想把 lora 权重合并到 base 模型上，可以通过离线的方式合并后得到权重文件。执行`merge_lora_weight.sh`即可完成 lora 权重和 base 模型权重的合并，其中的参数有：
+
+- `--base_model_path`：base 模型的权重目录
+- `--adapter_model_path`：lora 权重目录
+- `--output_path`：合并后的权重保存目录
+- `--save_dtype`： 以什么数据格式存储合并后的权重，可选值：fp16，bf16，fp32
+
+#### LLaMA-Factory 训练
+
+如果对 LLaMA-Factory 较为熟悉，可使用 LLaMA-Factory 进行微调。脚本、代码以及配置文件都归档在 `train/llama_factory_support` 目录下。如果没有特别说明，接下来我们提到的文件都是该目录下的文件。
+
+##### 安装
+
+可以通过下载源码 https://github.com/hiyouga/LLaMA-Factory/tree/main ，根据网站的指引进行安装。
+
+##### 训练脚本与配置文件
+
+各模型对应的配置文件和启动脚本如下：
+
+| 模型 | 全量微调配置 | LoRA 微调配置 | 启动脚本 |
+|------|------------|-------------|---------|
+| Hy-MT2-1.8B（Dense） | `hy_dense_1_8b_full_sft.yaml` | `hy_dense_1_8b_lora_sft.yaml` | `bash train_lf_dense.sh` |
+| Hy-MT2-7B（Dense） | `hy_dense_7b_full_sft.yaml` | `hy_dense_7b_lora_sft.yaml` | `YAML_FILE=hy_dense_7b_full_sft.yaml bash train_lf_dense.sh` |
+| Hy-MT2-30B-A3B（MoE） | `hy_v3_full_sft.yaml` | `hy_v3_lora_sft.yaml` | `bash train_lf.sh` |
+
+> **提示**：Dense 模型的启动脚本 `train_lf_dense.sh` 默认使用 `hy_dense_1_8b_full_sft.yaml`，可通过 `YAML_FILE` 环境变量指定其他配置文件。
+
+脚本中的关键参数如下：
+
+**模型相关：**
+
+- `model_name_or_path`: Hy-MT HF 格式预训练模型权重路径
+- `trust_remote_code`: 是否信任远程代码, Hy-MT 需要设置为 `true`
+
+**训练方法：**
+
+- `stage`: 训练阶段, 当前为 `sft`(监督微调)
+- `finetuning_type`: 微调类型, 可选 `full`(全量微调) 或 `lora`(LoRA 微调)
+- `deepspeed`: DeepSpeed 配置文件路径, 全量微调推荐 `ds_zero3_offload.json`, LoRA 微调推荐 `ds_zero2_offload_lora.json`
+
+**LoRA 参数(仅 LoRA 微调时生效)：**
+
+- `lora_rank`: LoRA 秩, 默认 `64`
+- `lora_alpha`: LoRA alpha 系数, 默认 `128`
+- `lora_dropout`: LoRA dropout 比率, 默认 `0.05`
+- `lora_target`: LoRA 应用的目标模块, 默认为 `q_proj,k_proj,v_proj,o_proj`
+
+**数据集：**
+
+- `dataset_dir`: 数据集目录路径
+- `dataset`: 数据集名称, 需要在 `dataset_dir` 下的 `dataset_info.json` 中注册
+- `template`: 对话模板, Hy-MT2-1.8B 使用 `hy_dense_1_8b`, Hy-MT2-7B 使用 `hy_dense_7b`, Hy-MT2-30B-A3B 使用 `hy_v3`
+- `cutoff_len`: 最大序列长度, 超出部分会被截断; 全量微调可设为 `262144`(262K), LoRA 微调建议设为 `8192` 以节省显存
+- `max_samples`: 每个数据集最多使用的样本数
+- `overwrite_cache`: 是否覆盖已缓存的预处理数据集
+
+**输出：**
+
+- `output_dir`: 输出目录, 日志、TensorBoard 和权重都会存储在此路径下
+- `logging_steps`: 每多少步记录一次日志
+- `save_steps`: 每多少步保存一次 checkpoint
+- `plot_loss`: 是否绘制训练 loss 曲线
+- `overwrite_output_dir`: 是否覆盖已有的输出目录
+- `save_only_model`: 是否只保存模型权重(不保存优化器状态等)
+- `report_to`: 日志上报工具, 可选 `none`, `wandb`, `tensorboard`, `swanlab`, `mlflow`
+
+**训练超参数：**
+
+- `per_device_train_batch_size`: 每张卡上的 batch size
+- `gradient_accumulation_steps`: 梯度累积步数, `per_device_train_batch_size * gradient_accumulation_steps * dp_size` 为 global batch size
+- `learning_rate`: 最大学习率, 全量微调推荐 `1.0e-5`, LoRA 微调推荐 `2.0e-4`
+- `num_train_epochs`: 训练轮数
+- `lr_scheduler_type`: 学习率调度器类型, 推荐使用 `cosine_with_min_lr`
+- `lr_scheduler_kwargs.min_lr_rate`: 最小学习率与最大学习率的比值, 例如 `0.1` 表示最小学习率为最大学习率的 10%
+- `warmup_ratio`: 预热阶段占总训练步数的比例
+- `bf16`: 是否使用 BFloat16 混合精度训练
+- `gradient_checkpointing`: 是否开启梯度重计算以节省显存
+- `ddp_timeout`: 分布式训练超时时间(毫秒)
+- `flash_attn`: 注意力实现方式, 推荐 `fa2`(FlashAttention-2), 也可选 `sdpa`; 使用 `fa2` 需要安装 flash-attn 包
+- `resume_from_checkpoint`: 从指定 checkpoint 路径恢复训练, 设为 `null` 表示从头开始训练
+
+##### 启动训练
+
+如需多机训练，请先完成 [配置机器间免密 ssh 登录](#配置机器间免密-ssh-登录多机训练) 中的配置（单机训练可跳过此步骤）。
+
+修改对应启动脚本中开头的以下配置：
+
+```shell
+export HOST_GPU_NUM=8
+# IP list, comma separated. e.g. "192.168.1.1,192.168.1.2" or single node "192.168.1.1"
+export IP_LIST=${IP_LIST:-"127.0.0.1"}
+```
+
+注意：如果`IP_LIST`环境变量未设置，则将`IP_LIST`替换为IP列表！格式为：
+```
+如果只有一个IP：
+IP_LIST=${ip_1}
+
+如果有多个IP：
+IP_LIST=${ip_1},${ip_2}
+
+```
+
+请将`${ip_1}`和`${ip_2}`替换为真实的IP地址。
+
+然后，在每一台机器上，在`train/llama_factory_support/`目录下执行对应的启动脚本即可。例如：
+
+```sh
+# Dense 1.8B 全量微调
+bash train_lf_dense.sh
+
+# Dense 7B LoRA 微调
+YAML_FILE=hy_dense_7b_lora_sft.yaml bash train_lf_dense.sh
+
+# MoE 30B-A3B 全量微调
+bash train_lf.sh
+```

train/data/example_data.jsonl

@@ -0,0 +1,100 @@
+{"messages": [{"role": "user", "content": "将以下中文翻译为英文，只输出翻译结果，不要额外解释：\n\n实验结果证明了假设的正确性。"}, {"role": "assistant", "content": "The experimental results demonstrate the correctness of the hypothesis."}]}
+{"messages": [{"role": "user", "content": "将以下中文翻译为英文，只输出翻译结果，不要额外解释：\n\n这篇论文提出了一个新颖的方法。"}, {"role": "assistant", "content": "This paper proposes a novel method."}]}
+{"messages": [{"role": "user", "content": "Translate the following English text into Chinese, output only the translation result without additional explanation:\n\nThis paper proposes a novel method."}, {"role": "assistant", "content": "这篇论文提出了一个新颖的方法。"}]}
+{"messages": [{"role": "user", "content": "将以下中文翻译为英文，只输出翻译结果，不要额外解释：\n\n谢谢你的帮助，我非常感激。"}, {"role": "assistant", "content": "Thank you for your help, I really appreciate it."}]}
+{"messages": [{"role": "user", "content": "Translate the following English text into Chinese, output only the translation result without additional explanation:\n\nPlease process this request as soon as possible."}, {"role": "assistant", "content": "请尽快处理这个请求。"}]}
+{"messages": [{"role": "user", "content": "Translate the following English text into Chinese, output only the translation result without additional explanation:\n\nHello! The weather is really nice today."}, {"role": "assistant", "content": "你好！今天天气真不错。"}]}
+{"messages": [{"role": "user", "content": "将以下中文翻译为英文，只输出翻译结果，不要额外解释：\n\n你最近在忙什么？"}, {"role": "assistant", "content": "What have you been up to lately?"}]}
+{"messages": [{"role": "user", "content": "Translate the following English text into Chinese, output only the translation result without additional explanation:\n\nThe global climate change conference was held this week."}, {"role": "assistant", "content": "全球气候变化会议在本周举行。"}]}
+{"messages": [{"role": "user", "content": "将以下中文翻译为英文，只输出翻译结果，不要额外解释：\n\n请尽快处理这个请求。"}, {"role": "assistant", "content": "Please process this request as soon as possible."}]}
+{"messages": [{"role": "user", "content": "Translate the following English text into Chinese, output only the translation result without additional explanation:\n\nHow do I get from the airport to the city center?"}, {"role": "assistant", "content": "从机场到市中心怎么走？"}]}
+{"messages": [{"role": "user", "content": "Translate the following English text into Chinese, output only the translation result without additional explanation:\n\nThe leaders of the two countries held talks and discussed bilateral relations."}, {"role": "assistant", "content": "两国领导人举行了会谈，讨论了双边关系。"}]}
+{"messages": [{"role": "user", "content": "将以下中文翻译为英文，只输出翻译结果，不要额外解释：\n\n请查收附件中的合同文件。"}, {"role": "assistant", "content": "Please find the contract document attached."}]}
+{"messages": [{"role": "user", "content": "Translate the following English text into Chinese, output only the translation result without additional explanation:\n\nI have an important meeting tomorrow."}, {"role": "assistant", "content": "我明天有个重要的会议。"}]}
+{"messages": [{"role": "user", "content": "Translate the following English text into Chinese, output only the translation result without additional explanation:\n\nDo you prefer Chinese food or Western food?"}, {"role": "assistant", "content": "你喜欢吃中餐还是西餐？"}]}
+{"messages": [{"role": "user", "content": "Translate the following English text into Chinese, output only the translation result without additional explanation:\n\nResearch has found that regular exercise helps improve mental health."}, {"role": "assistant", "content": "研究发现，定期锻炼有助于改善心理健康。"}]}
+{"messages": [{"role": "user", "content": "将以下中文翻译为英文，只输出翻译结果，不要额外解释：\n\n学习一门新语言需要时间和耐心。"}, {"role": "assistant", "content": "Learning a new language requires time and patience."}]}
+{"messages": [{"role": "user", "content": "将以下中文翻译为英文，只输出翻译结果，不要额外解释：\n\n这种药物的副作用是什么？"}, {"role": "assistant", "content": "What are the side effects of this medication?"}]}
+{"messages": [{"role": "user", "content": "Translate the following English text into Chinese, output only the translation result without additional explanation:\n\nThe experimental results demonstrate the correctness of the hypothesis."}, {"role": "assistant", "content": "实验结果证明了假设的正确性。"}]}
+{"messages": [{"role": "user", "content": "Translate the following English text into Chinese, output only the translation result without additional explanation:\n\nWhat is the landmark building of this city?"}, {"role": "assistant", "content": "这个城市的标志性建筑是什么？"}]}
+{"messages": [{"role": "user", "content": "将以下中文翻译为英文，只输出翻译结果，不要额外解释：\n\n这个API接口需要添加身份验证。"}, {"role": "assistant", "content": "This API endpoint needs to add authentication."}]}
+{"messages": [{"role": "user", "content": "将以下中文翻译为英文，只输出翻译结果，不要额外解释：\n\n门票多少钱一张？"}, {"role": "assistant", "content": "How much is one ticket?"}]}
+{"messages": [{"role": "user", "content": "将以下中文翻译为英文，只输出翻译结果，不要额外解释：\n\n该公司发布了最新款的智能手机。"}, {"role": "assistant", "content": "The company released its latest smartphone model."}]}
+{"messages": [{"role": "user", "content": "将以下中文翻译为英文，只输出翻译结果，不要额外解释：\n\n这里的当地特色美食是什么？"}, {"role": "assistant", "content": "What are the local specialty foods here?"}]}
+{"messages": [{"role": "user", "content": "Translate the following English text into Chinese, output only the translation result without additional explanation:\n\nIs the library also open on weekends?"}, {"role": "assistant", "content": "图书馆在周末也开放吗？"}]}
+{"messages": [{"role": "user", "content": "Translate the following English text into Chinese, output only the translation result without additional explanation:\n\nHow much is one ticket?"}, {"role": "assistant", "content": "门票多少钱一张？"}]}
+{"messages": [{"role": "user", "content": "将以下中文翻译为英文，只输出翻译结果，不要额外解释：\n\n我需要预约医生。"}, {"role": "assistant", "content": "I need to make an appointment with a doctor."}]}
+{"messages": [{"role": "user", "content": "Translate the following English text into Chinese, output only the translation result without additional explanation:\n\nThis program runs too slowly."}, {"role": "assistant", "content": "这个程序的运行速度太慢了。"}]}
+{"messages": [{"role": "user", "content": "将以下中文翻译为英文，只输出翻译结果，不要额外解释：\n\n这个城市的标志性建筑是什么？"}, {"role": "assistant", "content": "What is the landmark building of this city?"}]}
+{"messages": [{"role": "user", "content": "Translate the following English text into Chinese, output only the translation result without additional explanation:\n\nThis API endpoint needs to add authentication."}, {"role": "assistant", "content": "这个API接口需要添加身份验证。"}]}
+{"messages": [{"role": "user", "content": "将以下中文翻译为英文，只输出翻译结果，不要额外解释：\n\n这个程序的运行速度太慢了。"}, {"role": "assistant", "content": "This program runs too slowly."}]}
+{"messages": [{"role": "user", "content": "将以下中文翻译为英文，只输出翻译结果，不要额外解释：\n\n周末你有什么计划吗？"}, {"role": "assistant", "content": "Do you have any plans for the weekend?"}]}
+{"messages": [{"role": "user", "content": "将以下中文翻译为英文，只输出翻译结果，不要额外解释：\n\n你应该多休息，不要太劳累。"}, {"role": "assistant", "content": "You should rest more and not overwork yourself."}]}
+{"messages": [{"role": "user", "content": "Translate the following English text into Chinese, output only the translation result without additional explanation:\n\nThank you for your help, I really appreciate it."}, {"role": "assistant", "content": "谢谢你的帮助，我非常感激。"}]}
+{"messages": [{"role": "user", "content": "Translate the following English text into Chinese, output only the translation result without additional explanation:\n\nRegarding the previous order, I would like to confirm the delivery time."}, {"role": "assistant", "content": "关于上次的订单，我想确认一下发货时间。"}]}
+{"messages": [{"role": "user", "content": "Translate the following English text into Chinese, output only the translation result without additional explanation:\n\nThe new expressway will open to traffic next month."}, {"role": "assistant", "content": "新的高速公路将于下个月通车。"}]}
+{"messages": [{"role": "user", "content": "将以下中文翻译为英文，只输出翻译结果，不要额外解释：\n\n这个理论是由哪位科学家提出的？"}, {"role": "assistant", "content": "Which scientist proposed this theory?"}]}
+{"messages": [{"role": "user", "content": "Translate the following English text into Chinese, output only the translation result without additional explanation:\n\nWhat are the side effects of this medication?"}, {"role": "assistant", "content": "这种药物的副作用是什么？"}]}
+{"messages": [{"role": "user", "content": "Translate the following English text into Chinese, output only the translation result without additional explanation:\n\nI think this movie is absolutely wonderful."}, {"role": "assistant", "content": "我觉得这部电影非常精彩。"}]}
+{"messages": [{"role": "user", "content": "Translate the following English text into Chinese, output only the translation result without additional explanation:\n\nThe database server has encountered a failure."}, {"role": "assistant", "content": "数据库服务器出现了故障。"}]}
+{"messages": [{"role": "user", "content": "将以下中文翻译为英文，只输出翻译结果，不要额外解释：\n\n请将报价单发送给我。"}, {"role": "assistant", "content": "Please send me the quotation."}]}
+{"messages": [{"role": "user", "content": "将以下中文翻译为英文，只输出翻译结果，不要额外解释：\n\n请确保所有代码都通过了测试。"}, {"role": "assistant", "content": "Please ensure that all code has passed the tests."}]}
+{"messages": [{"role": "user", "content": "将以下中文翻译为英文，只输出翻译结果，不要额外解释：\n\n两国领导人举行了会谈，讨论了双边关系。"}, {"role": "assistant", "content": "The leaders of the two countries held talks and discussed bilateral relations."}]}
+{"messages": [{"role": "user", "content": "将以下中文翻译为英文，只输出翻译结果，不要额外解释：\n\n我们需要优化算法的性能。"}, {"role": "assistant", "content": "We need to optimize the performance of the algorithm."}]}
+{"messages": [{"role": "user", "content": "将以下中文翻译为英文，只输出翻译结果，不要额外解释：\n\n这个框架支持多种编程语言。"}, {"role": "assistant", "content": "This framework supports multiple programming languages."}]}
+{"messages": [{"role": "user", "content": "将以下中文翻译为英文，只输出翻译结果，不要额外解释：\n\n政府宣布了一系列新的经济政策。"}, {"role": "assistant", "content": "The government announced a series of new economic policies."}]}
+{"messages": [{"role": "user", "content": "Translate the following English text into Chinese, output only the translation result without additional explanation:\n\nOur products have passed international quality certification."}, {"role": "assistant", "content": "我们的产品通过了国际质量认证。"}]}
+{"messages": [{"role": "user", "content": "Translate the following English text into Chinese, output only the translation result without additional explanation:\n\nWe will have a meeting at 10 AM next Monday to discuss this project."}, {"role": "assistant", "content": "我们下周一上午10点开会讨论这个项目。"}]}
+{"messages": [{"role": "user", "content": "Translate the following English text into Chinese, output only the translation result without additional explanation:\n\nPlease find the contract document attached."}, {"role": "assistant", "content": "请查收附件中的合同文件。"}]}
+{"messages": [{"role": "user", "content": "将以下中文翻译为英文，只输出翻译结果，不要额外解释：\n\n这个问题我不太确定，让我查一下。"}, {"role": "assistant", "content": "I'm not quite sure about this issue, let me check."}]}
+{"messages": [{"role": "user", "content": "将以下中文翻译为英文，只输出翻译结果，不要额外解释：\n\n我们需要修复这个安全漏洞。"}, {"role": "assistant", "content": "We need to fix this security vulnerability."}]}
+{"messages": [{"role": "user", "content": "Translate the following English text into Chinese, output only the translation result without additional explanation:\n\nI need to make an appointment with a doctor."}, {"role": "assistant", "content": "我需要预约医生。"}]}
+{"messages": [{"role": "user", "content": "将以下中文翻译为英文，只输出翻译结果，不要额外解释：\n\n图书馆在周末也开放吗？"}, {"role": "assistant", "content": "Is the library also open on weekends?"}]}
+{"messages": [{"role": "user", "content": "将以下中文翻译为英文，只输出翻译结果，不要额外解释：\n\n我觉得这部电影非常精彩。"}, {"role": "assistant", "content": "I think this movie is absolutely wonderful."}]}
+{"messages": [{"role": "user", "content": "Translate the following English text into Chinese, output only the translation result without additional explanation:\n\nWhat have you been up to lately?"}, {"role": "assistant", "content": "你最近在忙什么？"}]}
+{"messages": [{"role": "user", "content": "将以下中文翻译为英文，只输出翻译结果，不要额外解释：\n\n我们的产品通过了国际质量认证。"}, {"role": "assistant", "content": "Our products have passed international quality certification."}]}
+{"messages": [{"role": "user", "content": "将以下中文翻译为英文，只输出翻译结果，不要额外解释：\n\n系统维护将在今晚凌晨2点进行。"}, {"role": "assistant", "content": "System maintenance will be performed at 2 AM tonight."}]}
+{"messages": [{"role": "user", "content": "Translate the following English text into Chinese, output only the translation result without additional explanation:\n\nDo you have any plans for the weekend?"}, {"role": "assistant", "content": "周末你有什么计划吗？"}]}
+{"messages": [{"role": "user", "content": "Translate the following English text into Chinese, output only the translation result without additional explanation:\n\nThe government announced a series of new economic policies."}, {"role": "assistant", "content": "政府宣布了一系列新的经济政策。"}]}
+{"messages": [{"role": "user", "content": "将以下中文翻译为英文，只输出翻译结果，不要额外解释：\n\n研究发现，定期锻炼有助于改善心理健康。"}, {"role": "assistant", "content": "Research has found that regular exercise helps improve mental health."}]}
+{"messages": [{"role": "user", "content": "Translate the following English text into Chinese, output only the translation result without additional explanation:\n\nPlease send me the quotation."}, {"role": "assistant", "content": "请将报价单发送给我。"}]}
+{"messages": [{"role": "user", "content": "将以下中文翻译为英文，只输出翻译结果，不要额外解释：\n\n新的高速公路将于下个月通车。"}, {"role": "assistant", "content": "The new expressway will open to traffic next month."}]}
+{"messages": [{"role": "user", "content": "将以下中文翻译为英文，只输出翻译结果，不要额外解释：\n\n我明天有个重要的会议。"}, {"role": "assistant", "content": "I have an important meeting tomorrow."}]}
+{"messages": [{"role": "user", "content": "将以下中文翻译为英文，只输出翻译结果，不要额外解释：\n\n你喜欢吃中餐还是西餐？"}, {"role": "assistant", "content": "Do you prefer Chinese food or Western food?"}]}
+{"messages": [{"role": "user", "content": "Translate the following English text into Chinese, output only the translation result without additional explanation:\n\nCould you provide a detailed technical proposal?"}, {"role": "assistant", "content": "能否提供一个详细的技术方案？"}]}
+{"messages": [{"role": "user", "content": "Translate the following English text into Chinese, output only the translation result without additional explanation:\n\nWhat are the local specialty foods here?"}, {"role": "assistant", "content": "这里的当地特色美食是什么？"}]}
+{"messages": [{"role": "user", "content": "将以下中文翻译为英文，只输出翻译结果，不要额外解释：\n\n科学家在量子计算领域取得了重大突破。"}, {"role": "assistant", "content": "Scientists have made a major breakthrough in the field of quantum computing."}]}
+{"messages": [{"role": "user", "content": "Translate the following English text into Chinese, output only the translation result without additional explanation:\n\nAre there any good restaurants you would recommend nearby?"}, {"role": "assistant", "content": "这附近有什么好吃的餐厅推荐吗？"}]}
+{"messages": [{"role": "user", "content": "将以下中文翻译为英文，只输出翻译结果，不要额外解释：\n\n请在下周五之前提交你的研究报告。"}, {"role": "assistant", "content": "Please submit your research report by next Friday."}]}
+{"messages": [{"role": "user", "content": "Translate the following English text into Chinese, output only the translation result without additional explanation:\n\nThis framework supports multiple programming languages."}, {"role": "assistant", "content": "这个框架支持多种编程语言。"}]}
+{"messages": [{"role": "user", "content": "将以下中文翻译为英文，只输出翻译结果，不要额外解释：\n\n全球气候变化会议在本周举行。"}, {"role": "assistant", "content": "The global climate change conference was held this week."}]}
+{"messages": [{"role": "user", "content": "将以下中文翻译为英文，只输出翻译结果，不要额外解释：\n\n你好！今天天气真不错。"}, {"role": "assistant", "content": "Hello! The weather is really nice today."}]}
+{"messages": [{"role": "user", "content": "Translate the following English text into Chinese, output only the translation result without additional explanation:\n\nLearning a new language requires time and patience."}, {"role": "assistant", "content": "学习一门新语言需要时间和耐心。"}]}
+{"messages": [{"role": "user", "content": "Translate the following English text into Chinese, output only the translation result without additional explanation:\n\nI'm not quite sure about this issue, let me check."}, {"role": "assistant", "content": "这个问题我不太确定，让我查一下。"}]}
+{"messages": [{"role": "user", "content": "Translate the following English text into Chinese, output only the translation result without additional explanation:\n\nPlease submit your research report by next Friday."}, {"role": "assistant", "content": "请在下周五之前提交你的研究报告。"}]}
+{"messages": [{"role": "user", "content": "Translate the following English text into Chinese, output only the translation result without additional explanation:\n\nCould you do me a favor?"}, {"role": "assistant", "content": "你能帮我一个忙吗？"}]}
+{"messages": [{"role": "user", "content": "Translate the following English text into Chinese, output only the translation result without additional explanation:\n\nWe need to optimize the performance of the algorithm."}, {"role": "assistant", "content": "我们需要优化算法的性能。"}]}
+{"messages": [{"role": "user", "content": "将以下中文翻译为英文，只输出翻译结果，不要额外解释：\n\n人工智能技术正在改变各行各业。"}, {"role": "assistant", "content": "Artificial intelligence technology is transforming various industries."}]}
+{"messages": [{"role": "user", "content": "Translate the following English text into Chinese, output only the translation result without additional explanation:\n\nThe company released its latest smartphone model."}, {"role": "assistant", "content": "该公司发布了最新款的智能手机。"}]}
+{"messages": [{"role": "user", "content": "Translate the following English text into Chinese, output only the translation result without additional explanation:\n\nWe need to fix this security vulnerability."}, {"role": "assistant", "content": "我们需要修复这个安全漏洞。"}]}
+{"messages": [{"role": "user", "content": "Translate the following English text into Chinese, output only the translation result without additional explanation:\n\nI would like to book a double room for two nights."}, {"role": "assistant", "content": "我想预订一间双人房，住两晚。"}]}
+{"messages": [{"role": "user", "content": "Translate the following English text into Chinese, output only the translation result without additional explanation:\n\nThank you for your cooperation, and we look forward to more opportunities in the future."}, {"role": "assistant", "content": "感谢贵公司的合作，期待未来更多机会。"}]}
+{"messages": [{"role": "user", "content": "将以下中文翻译为英文，只输出翻译结果，不要额外解释：\n\n这附近有什么好吃的餐厅推荐吗？"}, {"role": "assistant", "content": "Are there any good restaurants you would recommend nearby?"}]}
+{"messages": [{"role": "user", "content": "将以下中文翻译为英文，只输出翻译结果，不要额外解释：\n\n我想预订一间双人房，住两晚。"}, {"role": "assistant", "content": "I would like to book a double room for two nights."}]}
+{"messages": [{"role": "user", "content": "Translate the following English text into Chinese, output only the translation result without additional explanation:\n\nSystem maintenance will be performed at 2 AM tonight."}, {"role": "assistant", "content": "系统维护将在今晚凌晨2点进行。"}]}
+{"messages": [{"role": "user", "content": "Translate the following English text into Chinese, output only the translation result without additional explanation:\n\nYou should rest more and not overwork yourself."}, {"role": "assistant", "content": "你应该多休息，不要太劳累。"}]}
+{"messages": [{"role": "user", "content": "将以下中文翻译为英文，只输出翻译结果，不要额外解释：\n\n这个地方真美，我们应该经常来。"}, {"role": "assistant", "content": "This place is so beautiful, we should come here more often."}]}
+{"messages": [{"role": "user", "content": "Translate the following English text into Chinese, output only the translation result without additional explanation:\n\nThis place is so beautiful, we should come here more often."}, {"role": "assistant", "content": "这个地方真美，我们应该经常来。"}]}
+{"messages": [{"role": "user", "content": "Translate the following English text into Chinese, output only the translation result without additional explanation:\n\nArtificial intelligence technology is transforming various industries."}, {"role": "assistant", "content": "人工智能技术正在改变各行各业。"}]}
+{"messages": [{"role": "user", "content": "将以下中文翻译为英文，只输出翻译结果，不要额外解释：\n\n我们下周一上午10点开会讨论这个项目。"}, {"role": "assistant", "content": "We will have a meeting at 10 AM next Monday to discuss this project."}]}
+{"messages": [{"role": "user", "content": "Translate the following English text into Chinese, output only the translation result without additional explanation:\n\nScientists have made a major breakthrough in the field of quantum computing."}, {"role": "assistant", "content": "科学家在量子计算领域取得了重大突破。"}]}
+{"messages": [{"role": "user", "content": "Translate the following English text into Chinese, output only the translation result without additional explanation:\n\nPlease ensure that all code has passed the tests."}, {"role": "assistant", "content": "请确保所有代码都通过了测试。"}]}
+{"messages": [{"role": "user", "content": "将以下中文翻译为英文，只输出翻译结果，不要额外解释：\n\n关于上次的订单，我想确认���下发货时间。"}, {"role": "assistant", "content": "Regarding the previous order, I would like to confirm the delivery time."}]}
+{"messages": [{"role": "user", "content": "将以下中文翻译为英文，只输出翻译结果，不要额外解释：\n\n能否提供一个详细的技术方案？"}, {"role": "assistant", "content": "Could you provide a detailed technical proposal?"}]}
+{"messages": [{"role": "user", "content": "将以下中文翻译为英文，只输出翻译结果，不要额外解释：\n\n从机场到市中心怎么走？"}, {"role": "assistant", "content": "How do I get from the airport to the city center?"}]}
+{"messages": [{"role": "user", "content": "将以下中文翻译为英文，只输出翻译结果，不要额外解释：\n\n我需要办理签证才能去那个国家吗？"}, {"role": "assistant", "content": "Do I need a visa to go to that country?"}]}
+{"messages": [{"role": "user", "content": "将以下中文翻译为英文，只输出翻译结果，不要额外解释：\n\n数据库服务器出现了故障。"}, {"role": "assistant", "content": "The database server has encountered a failure."}]}
+{"messages": [{"role": "user", "content": "Translate the following English text into Chinese, output only the translation result without additional explanation:\n\nWhich scientist proposed this theory?"}, {"role": "assistant", "content": "这个理论是由哪位科学家提出的？"}]}
+{"messages": [{"role": "user", "content": "将以下中文翻译为英文，只输出翻译结果，不要额外解释：\n\n你能帮我一个忙吗？"}, {"role": "assistant", "content": "Could you do me a favor?"}]}
+{"messages": [{"role": "user", "content": "将以下中文翻译为英文，只输出翻译结果，不要额外解释：\n\n感谢贵公司的合作，期待未来更多机会。"}, {"role": "assistant", "content": "Thank you for your cooperation, and we look forward to more opportunities in the future."}]}
+{"messages": [{"role": "user", "content": "Translate the following English text into Chinese, output only the translation result without additional explanation:\n\nDo I need a visa to go to that country?"}, {"role": "assistant", "content": "我需要办理签证才能去那个国家吗？"}]}

train/deepspeed_support/ds_zero2_no_offload.json

@@ -0,0 +1,20 @@
+{
+    "fp16": {
+        "enabled": false
+    },
+    "zero_optimization": {
+        "stage": 2,
+        "allgather_partitions": true,
+        "allgather_bucket_size": 1e8,
+        "overlap_comm": true,
+        "reduce_scatter": true,
+        "reduce_bucket_size": 1e8,
+        "contiguous_gradients": true
+    },
+    "gradient_accumulation_steps": "auto",
+    "gradient_clipping": "auto",
+    "steps_per_print": 10,
+    "train_batch_size": "auto",
+    "train_micro_batch_size_per_gpu": "auto",
+    "wall_clock_breakdown": false
+}

train/deepspeed_support/ds_zero2_offload.json

@@ -0,0 +1,27 @@
+{
+    "fp16": {
+        "enabled": false
+    },
+    "bf16": {
+        "enabled": true
+    },
+    "zero_optimization": {
+        "stage": 2,
+        "offload_optimizer": {
+            "device": "cpu",
+            "pin_memory": true
+        },
+        "allgather_partitions": true,
+        "allgather_bucket_size": 1e8,
+        "overlap_comm": true,
+        "reduce_scatter": true,
+        "reduce_bucket_size": 1e8,
+        "contiguous_gradients": true
+    },
+    "gradient_accumulation_steps": "auto",
+    "gradient_clipping": "auto",
+    "steps_per_print": 10,
+    "train_batch_size": "auto",
+    "train_micro_batch_size_per_gpu": "auto",
+    "wall_clock_breakdown": false
+}

train/deepspeed_support/ds_zero3_no_offload.json

@@ -0,0 +1,41 @@
+{
+    "fp16": {
+        "enabled": "auto",
+        "loss_scale": 0,
+        "loss_scale_window": 1000,
+        "initial_scale_power": 16,
+        "hysteresis": 2,
+        "min_loss_scale": 1
+    },
+    "bf16": {
+        "enabled": "auto"
+    },
+
+    "zero_optimization": {
+        "stage": 3,
+        "offload_optimizer": {
+            "device": "none",
+            "pin_memory": true
+        },
+        "offload_param": {
+            "device": "none",
+            "pin_memory": true
+        },
+        "overlap_comm": true,
+        "contiguous_gradients": true,
+        "sub_group_size": 1e9,
+        "reduce_bucket_size": "auto",
+        "stage3_prefetch_bucket_size": "auto",
+        "stage3_param_persistence_threshold": "auto",
+        "stage3_max_live_parameters": 1e9,
+        "stage3_max_reuse_distance": 1e9,
+        "stage3_gather_16bit_weights_on_model_save": true
+    },
+
+    "gradient_accumulation_steps": "auto",
+    "gradient_clipping": "auto",
+    "steps_per_print": 10,
+    "train_batch_size": "auto",
+    "train_micro_batch_size_per_gpu": "auto",
+    "wall_clock_breakdown": false
+}

train/deepspeed_support/ds_zero3_offload.json

@@ -0,0 +1,41 @@
+{
+    "fp16": {
+        "enabled": "auto",
+        "loss_scale": 0,
+        "loss_scale_window": 1000,
+        "initial_scale_power": 16,
+        "hysteresis": 2,
+        "min_loss_scale": 1
+    },
+    "bf16": {
+        "enabled": "auto"
+    },
+
+    "zero_optimization": {
+        "stage": 3,
+        "offload_optimizer": {
+            "device": "cpu",
+            "pin_memory": true
+        },
+        "offload_param": {
+            "device": "cpu",
+            "pin_memory": true
+        },
+        "overlap_comm": true,
+        "contiguous_gradients": true,
+        "sub_group_size": 1e9,
+        "reduce_bucket_size": "auto",
+        "stage3_prefetch_bucket_size": "auto",
+        "stage3_param_persistence_threshold": "auto",
+        "stage3_max_live_parameters": 1e9,
+        "stage3_max_reuse_distance": 1e9,
+        "stage3_gather_16bit_weights_on_model_save": false
+    },
+
+    "gradient_accumulation_steps": "auto",
+    "gradient_clipping": "auto",
+    "steps_per_print": 10,
+    "train_batch_size": "auto",
+    "train_micro_batch_size_per_gpu": "auto",
+    "wall_clock_breakdown": false
+}

train/deepspeed_support/ds_zero3_offload_no_auto.json

@@ -0,0 +1,36 @@
+{
+    "fp16": {
+        "enabled": false,
+        "loss_scale": 0,
+        "loss_scale_window": 1000,
+        "initial_scale_power": 16,
+        "hysteresis": 2,
+        "min_loss_scale": 1
+    },
+    "bf16": {
+        "enabled": true
+    },
+
+    "zero_optimization": {
+        "stage": 3,
+        "offload_optimizer": {
+            "device": "cpu",
+            "pin_memory": true
+        },
+        "overlap_comm": true,
+        "contiguous_gradients": true,
+        "sub_group_size": 1e9,
+        "reduce_bucket_size": 1e8,
+        "stage3_prefetch_bucket_size": 1e8,
+        "stage3_param_persistence_threshold": 0,
+        "stage3_max_live_parameters": 1e8,
+        "stage3_max_reuse_distance": 1e8,
+        "stage3_gather_16bit_weights_on_model_save": true
+    },
+
+    "gradient_accumulation_steps": 1,
+    "gradient_clipping": "auto",
+    "steps_per_print": 10,
+    "train_micro_batch_size_per_gpu": 1,
+    "wall_clock_breakdown": false
+}

train/deepspeed_support/merge_lora_weight.py

@@ -0,0 +1,54 @@
+# 导入所需的库
+from transformers import AutoModelForCausalLM  # 用于加载预训练的语言模型
+from peft import LoraConfig, get_peft_model, PeftModel  # 用于处理LoRA权重
+import argparse  # 用于解析命令行参数
+import shutil  # 用于文件操作，如复制
+import os  # 用于文件路径操作
+import torch  # 用于深度学习操作
+
+def main():
+    # 创建参数解析器
+    parser = argparse.ArgumentParser()
+    # 添加命令行参数
+    parser.add_argument("--base_model_path", type=str, required=True, 
+                        help="Path to pretrained model or model identifier from huggingface.co/models")
+    parser.add_argument("--adapter_model_path", type=str, required=True, help="Path to adapter model")
+    parser.add_argument("--output_path", type=str, required=True, help="Path to save the output model")
+    parser.add_argument("--save_dtype", type=str, choices=['bf16', 'fp32', 'fp16'], 
+                        default='fp32', help="In which dtype to save, fp32, bf16 or fp16.")
+    # 解析命令行参数
+    args = parser.parse_args()
+
+    name2dtype = {'bf16': torch.bfloat16, 'fp32': torch.float32, 'fp16': torch.float16}
+    # 加载基座模型
+    model = AutoModelForCausalLM.from_pretrained(
+        args.base_model_path, device_map='cpu', 
+        trust_remote_code=True, torch_dtype=name2dtype[args.save_dtype]
+    )
+    # 在基座模型的基础上加载 adapter 权重
+    model = PeftModel.from_pretrained(model, args.adapter_model_path, trust_remote_code=True)
+    # 融合模型和 adapter
+    model = model.merge_and_unload()
+    # 保存融合后的模型权重
+    model.save_pretrained(args.output_path, safe_serialization=False)
+
+    # Copy tokenizer, config and other non-weight files from base model
+    # Skip model weight files (.safetensors, .bin, .pt) and index files
+    _SKIP_SUFFIXES = ('.safetensors', '.bin', '.pt', '.pth')
+    _SKIP_NAMES = {'model.safetensors.index.json', 'pytorch_model.bin.index.json'}
+
+    for fname in os.listdir(args.base_model_path):
+        src = os.path.join(args.base_model_path, fname)
+        if not os.path.isfile(src):
+            continue
+        if fname in _SKIP_NAMES or fname.endswith(_SKIP_SUFFIXES):
+            continue
+        dst = os.path.join(args.output_path, fname)
+        if not os.path.exists(dst):
+            shutil.copy(src, dst)
+            print(f'Copied {fname}')
+
+    print(f'Merged model weight is saved to {args.output_path}')
+    
+if __name__ == "__main__":
+    main()

train/deepspeed_support/merge_lora_weight.sh

@@ -0,0 +1,6 @@
+SCRIPT_DIR=$(dirname "$0")
+python3 ${SCRIPT_DIR}/merge_lora_weight.py \
+    --base_model_path /xxx/hy_l_train/checkpoint-200 \
+    --adapter_model_path /xxx/runs/hy_l_lora_train/checkpoint-200 \
+    --output_path /xxx/ckpts/merged_hy_lora_weight \
+    --save_dtype bf16

train/deepspeed_support/train.py

@@ -0,0 +1,564 @@
+# Copyright 2024 Tencent Inc. 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.
+
+# Copyright 2022 EleutherAI and the HuggingFace Inc. team. All rights reserved.
+#
+# This code is based on EleutherAI's GPT-NeoX library and the GPT-NeoX
+# and OPT implementations in this library. It has been modified from its
+# original forms to accommodate minor architectural differences compared
+# to GPT-NeoX and OPT used by the Meta AI team that trained the model.
+#
+# 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 os
+import re
+import sys
+sys.path.insert(0, os.path.dirname(os.path.dirname(os.path.abspath(__file__))))
+import json
+import torch
+import shutil
+import logging
+from dataclasses import dataclass, field
+import deepspeed
+from typing import Optional, Dict
+
+import transformers
+from torch.utils.data import Dataset
+from transformers import Trainer, TrainerCallback
+from peft import LoraConfig, get_peft_model, PeftModel
+from transformers.trainer_utils import PREFIX_CHECKPOINT_DIR
+from transformers.modeling_utils import unwrap_model
+
+
+def print_args(args, name='arguments'):
+    """Print arguments."""
+    if torch.distributed.get_rank() == 0:
+        print(f'------------------------ {name} ------------------------', flush=True)
+        str_list = []
+        for arg in vars(args):
+            dots = '.' * (48 - len(arg))
+            str_list.append('  {} {} {}'.format(arg, dots, getattr(args, arg)))
+        for arg in sorted(str_list, key=lambda x: x.lower()):
+            print(arg, flush=True)
+        print(f'-------------------- end of {name} ---------------------', flush=True)
+
+
+@dataclass
+class ModelArguments:
+    use_flash_attn: bool = field(
+        default=False, 
+        metadata={"help": "Enable FlashAttention-2 for faster training."}
+    )
+    use_lora: bool = field(default=False, metadata={"help": "Enable Lora for faster training."})
+    hidden_size: int = field(default=2048, metadata={"help": "The hidden size of the model."})
+    num_layers: int = field(default=24, metadata={"help": "The number of layers of the model."})
+    num_attention_heads: int = field(default=16, metadata={"help": "The number of attention heads of the model."})
+    intermediate_size: int = field(default=8192, metadata={"help": "The intermediate size of the model."})
+    max_position_embeddings: int = field(
+        default=2048, 
+        metadata={"help": "The maximum sequence length that this model might ever be used with."}
+    )
+    vocab_size: int = field(default=50257, metadata={"help": "The vocabulary size of the model."})
+    type_vocab_size: int = field(default=1, metadata={"help": "The vocabulary size of the model."})
+    layer_norm_eps: float = field(
+        default=1e-5, 
+        metadata={"help": "The epsilon used by the layer normalization layers of the model."}
+    )
+    moe_topk: int = field(default=4, metadata={"help": "The topk for MOE."})
+    num_experts: int = field(default=8, metadata={"help": "The number of experts for MOE."})
+    num_key_value_heads: int = field(default=16, metadata={"help": "The number of key-value heads in GQA."})
+    moe_intermediate_size: int = field(default=1536, metadata={"help": "The intermediate size of each MoE expert."})
+    use_mixed_mlp_moe: bool = field(
+        default=False, 
+        metadata={"help": "Whether to use mixed MoE with shared expert."}
+    )
+    num_shared_expert: int = field(default=1, metadata={"help": "Number of shared experts."})
+    use_qk_norm: bool = field(default=False, metadata={"help": "Whether to use qk norm."})
+    moe_layer_num_skipped: int = field(default=1, metadata={"help": "Number of initial dense layers before MoE layers."})
+    tie_word_embeddings: bool = field(
+        default=True, 
+        metadata={"help": "Whether to tie the word embeddings of the encoder and the decoder."}
+    )
+    lora_rank: int = field(default=64, metadata={"help": "The rank of lora."})
+    lora_alpha: int = field(default=8, metadata={"help": "Lora alpha"})
+    lora_dropout: float = field(default=0.0, metadata={"help": "Lora dropout"})
+    train_attention_params_only: bool = field(default=False, metadata={
+        "help": "Whether to train attention parameters only."}
+    )
+
+
+@dataclass
+class DataArguments:
+    train_data_file: str = field(default=None, metadata={"help": "Path to the training data."})
+    max_seq_length: int = field(
+        default=2048, 
+        metadata={"help": "The max sequence length of the model inputs after tokenization."}
+    )
+    complex_data: Optional[str] = field(default=None)
+    use_dummy_data: bool = field(default=False, metadata={"help": "Use dummy data."})
+
+
+@dataclass
+class TrainingArguments(transformers.TrainingArguments):
+    cache_dir: Optional[str] = field(default=None)
+    optim: str = field(default="adamw_torch")
+    model_max_length: int = field(
+        default=2048,
+        metadata={"help": "Maximum sequence length. Sequences will be right padded (and possibly truncated)."},
+    )
+    tokenizer_name_or_path: Optional[str] = field(default=None)
+    model_name_or_path: Optional[str] = field(default=None)
+    min_lr: float = field(
+        default=0.01, 
+        metadata={"help": "The final learning rate at the end of the decay will be learning_rate * min_lr"}
+    )
+
+
+IGNORE_INDEX = -100
+
+
+class DummyDataset(Dataset):
+    def __init__(self, tokenizer, max_seq_length=512, length=1000):
+        self.tokenizer = tokenizer
+        self.max_seq_length = max_seq_length
+        self.length = length
+    
+    def __len__(self):
+        return self.length
+    
+    def __getitem__(self, index):
+        tokens = torch.randint(0, self.tokenizer.vocab_size, (self.max_seq_length, ))
+        return {'input_ids': tokens, 'labels': tokens}
+
+
+class SFTDataset(Dataset):
+    def __init__(self, data_file, tokenizer, max_seq_length = 2048, prompt_format = 'mplus'):
+        self.tokenizer = tokenizer
+        self.prompt_format = prompt_format
+        self.max_seq_length = max_seq_length
+
+        self.data_list = self.load_data(data_file)
+
+    def __len__(self):
+        return len(self.data_list)
+
+    def load_data(self, data_file):
+        logging.info('Loading data: {}'.format(data_file))
+        with open(data_file, 'r', encoding='utf8') as f:
+            data_list = f.readlines()
+        logging.info("there are {} data in dataset".format(len(data_list)))
+        return data_list
+
+    def encode_data(self, data_dict):
+        model_inputs = {}
+        reasoning_effort = data_dict.get('reasoning_effort', None)
+        if reasoning_effort is None:
+            reasoning_effort = 'no_think'
+        template_output = self.tokenizer.apply_chat_template(data_dict['messages'], tokenize=True, return_dict=False, is_training=True, reasoning_effort=reasoning_effort)
+        if isinstance(template_output, list) and len(template_output) > 0 and isinstance(template_output[0], list):
+            template_output = template_output[0]
+        message_tokens = torch.tensor(template_output, dtype=torch.long)
+
+        # Use new HunYuan tokenizer special tokens
+        assistant_token_id = self.tokenizer.convert_tokens_to_ids('<｜hy_Assistant｜>')
+        eos_token_id = self.tokenizer.convert_tokens_to_ids(self.tokenizer.eos_token)
+        pad_token_id = self.tokenizer.pad_token_id
+
+        # Find assistant reply boundaries: starts at <｜hy_Assistant｜>, ends at eos_token
+        loss_token_begins = (message_tokens == assistant_token_id).nonzero(as_tuple=True)[0].tolist()
+        loss_token_ends = (message_tokens == eos_token_id).nonzero(as_tuple=True)[0].tolist()
+        message_labels = torch.tensor([IGNORE_INDEX] * message_tokens.shape[0])
+        for begin_idx, end_idx in zip(loss_token_begins, loss_token_ends):
+            # Compute loss from the token after <｜hy_Assistant｜> to eos_token (inclusive)
+            message_labels[begin_idx + 1:end_idx + 1] = message_tokens[begin_idx + 1:end_idx + 1]
+        input_ids = message_tokens.to(torch.long)
+        labels = message_labels.to(torch.long)
+
+        input_ids = input_ids[:self.max_seq_length]
+        labels = labels[:self.max_seq_length]
+        attention_mask = [1 if val != pad_token_id else 0 for val in input_ids]
+        model_inputs["input_ids"] = input_ids
+        model_inputs["attention_mask"] = torch.tensor(attention_mask, dtype=torch.bool)
+        model_inputs["labels"] = labels
+
+        return model_inputs
+
+    def __getitem__(self, index):
+        data = self.data_list[index]
+        data = json.loads(data)
+        model_inputs = self.encode_data(data)
+
+        return model_inputs
+
+
+@dataclass
+class DataCollatorForSupervisedDataset(object):
+    """Collate examples for supervised fine-tuning."""
+
+    tokenizer: transformers.PreTrainedTokenizer
+
+    def __call__(self, instances):
+        input_ids = [instance['input_ids'] for instance in instances]
+        labels = [instance['labels'] for instance in instances]
+        pad_token_id = self.tokenizer.pad_token_id
+        input_ids = torch.nn.utils.rnn.pad_sequence(input_ids, batch_first=True, padding_value=pad_token_id)
+        labels = torch.nn.utils.rnn.pad_sequence(labels, batch_first=True, padding_value=IGNORE_INDEX)
+        return dict(
+            input_ids=input_ids,
+            labels=labels,
+            attention_mask=input_ids.ne(pad_token_id),
+        )
+
+
+def make_supervised_data_module(tokenizer, data_args) -> Dict:
+    """Make dataset and collator for supervised fine-tuning."""
+    if data_args.use_dummy_data:
+        train_dataset = DummyDataset(tokenizer, data_args.max_seq_length)
+    else:
+        train_dataset = SFTDataset(
+            tokenizer=tokenizer, 
+            data_file=data_args.train_data_file, 
+            max_seq_length=data_args.max_seq_length
+        )
+    data_collator = DataCollatorForSupervisedDataset(tokenizer=tokenizer)
+    return dict(train_dataset=train_dataset, eval_dataset=None, data_collator=data_collator)
+
+
+# for full model training, change the config.json, copy the model and configuration to support Auto load
+class CustomSaveCallback(TrainerCallback):
+    def on_save(self, args, state, control, **kwargs):
+        if torch.distributed.get_rank() == 0:
+            output_dir = os.path.join(args.output_dir, f"{PREFIX_CHECKPOINT_DIR}-{state.global_step}")
+
+            # Copy tokenizer files to checkpoint directory
+            tokenizer_files = [
+                'generation_config.json',
+                'hy.tiktoken',
+                'tokenizer_config.json',
+                'tokenization_hy.py',
+                'tokenizer.json',
+                'special_tokens_map.json',
+                'chat_template.jinja',
+            ]
+            for fname in tokenizer_files:
+                src = os.path.join(args.tokenizer_name_or_path, fname)
+                if os.path.isfile(src):
+                    shutil.copy(src, os.path.join(output_dir, fname))
+
+        return control
+
+
+def train():
+    parser = transformers.HfArgumentParser((ModelArguments, DataArguments, TrainingArguments))
+    model_args, data_args, training_args = parser.parse_args_into_dataclasses()
+    print_args(model_args, 'model arguments')
+    print_args(data_args, 'data arguments')
+    print_args(training_args, 'training arguments')
+
+    tokenizer = transformers.AutoTokenizer.from_pretrained(
+        training_args.tokenizer_name_or_path,
+        trust_remote_code = True
+    )
+
+    init_kwargs = {}
+    if model_args.use_flash_attn:
+        init_kwargs["attn_implementation"] = "flash_attention_2"
+        # Workaround: transformers >= 5.x uses importlib.metadata.packages_distributions()
+        # to verify flash-attn package name, which fails when the package is installed under
+        # a custom distribution name (e.g. ptm-flash-attn). Patch the check to skip it.
+        try:
+            from transformers.modeling_flash_attention_utils import FLASH_ATTENTION_COMPATIBILITY_MATRIX
+            _orig_pkg_check = FLASH_ATTENTION_COMPATIBILITY_MATRIX[2]["pkg_availability_check"]
+            FLASH_ATTENTION_COMPATIBILITY_MATRIX[2]["pkg_availability_check"] = lambda *a, **kw: True
+            print("[Patch] Bypassed flash_attn package distribution name check for FA2.")
+        except Exception as e:
+            print(f"[Patch] Could not patch FA2 pkg check (non-fatal): {e}")
+    if training_args.bf16:
+        init_kwargs["dtype"] = torch.bfloat16
+    elif training_args.fp16:
+        init_kwargs["dtype"] = torch.float16
+
+    # Check if model weights exist (not just the directory)
+    _has_weights = (
+        training_args.model_name_or_path is not None
+        and os.path.isdir(training_args.model_name_or_path)
+        and any(
+            os.path.isfile(os.path.join(training_args.model_name_or_path, f))
+            for f in ("model.safetensors", "pytorch_model.bin", "model.safetensors.index.json", "pytorch_model.bin.index.json")
+        )
+    )
+
+    # -----------------------------------------------------------------------
+    # Fix: Rename checkpoint keys so that old-style weight names (e.g.
+    # self_attn.q_norm) are mapped to the current model attribute names
+    # (e.g. self_attn.query_layernorm).  The model's
+    # _fix_state_dict_key_on_load hook is NOT invoked on the DeepSpeed
+    # ZeRO-3 loading path, so we monkey-patch the ZeRO-3 loader instead.
+    # -----------------------------------------------------------------------
+    _CKPT_KEY_RENAMES = [
+        ("mlp.gate.wg.", "mlp.router.gate."),
+    ]
+
+    from transformers.integrations.deepspeed import (
+        _load_state_dict_into_zero3_model as _orig_load_zero3,
+    )
+    import transformers.integrations.deepspeed as _ds_mod
+    import transformers.modeling_utils as _mu_mod
+
+    def _patched_load_zero3(model_to_load, state_dict, load_config=None):
+        new_sd = {}
+        for k, v in state_dict.items():
+            new_k = k
+            for old_sub, new_sub in _CKPT_KEY_RENAMES:
+                if old_sub in new_k:
+                    new_k = new_k.replace(old_sub, new_sub)
+                    break
+            new_sd[new_k] = v
+
+        # Call original ZeRO-3 loader for parameters
+        result = _orig_load_zero3(model_to_load, new_sd, load_config)
+
+        # -------------------------------------------------------------------
+        # Patch: Manually load buffers (e.g. e_score_correction_bias).
+        # ZeRO-3's loader only handles named_parameters, not named_buffers.
+        # -------------------------------------------------------------------
+        buffers_loaded = 0
+        for name, buf in model_to_load.named_buffers():
+            if name in new_sd:
+                src_tensor = new_sd[name]
+                if isinstance(src_tensor, torch.Tensor):
+                    buf.data.copy_(src_tensor.to(buf.dtype))
+                    buffers_loaded += 1
+                    # Remove from unexpected keys if tracked
+                    if isinstance(result, tuple) and len(result) >= 2:
+                        if isinstance(result[1], set):
+                            result[1].discard(name)
+        if buffers_loaded > 0:
+            print(f"[HYV3 Patch] Manually loaded {buffers_loaded} buffers "
+                  f"(e.g. e_score_correction_bias) into model.")
+
+        return result
+
+    _ds_mod._load_state_dict_into_zero3_model = _patched_load_zero3
+    _mu_mod._load_state_dict_into_zero3_model = _patched_load_zero3
+    # -----------------------------------------------------------------------
+
+    # -------------------------------------------------------------------
+    # Patch: Save-time reverse key rename + 3D -> per-expert unfuse.
+    #
+    # When saving checkpoints, the model state_dict uses 3D fused experts
+    # and new naming.  We reverse both for old checkpoint compatibility:
+    #   - mlp.gate.           -> mlp.router.gate.
+    #   - mlp.e_score_correction_bias -> mlp.expert_bias
+    #   - mlp.shared_experts. -> mlp.shared_mlp.
+    #   - experts.gate_up_proj -> experts.{N}.gate_proj.weight + up_proj
+    #   - experts.down_proj    -> experts.{N}.down_proj.weight
+    # -------------------------------------------------------------------
+    _SAVE_KEY_RENAMES = [
+        ("mlp.gate.", "mlp.router.gate."),
+        ("mlp.e_score_correction_bias", "mlp.expert_bias"),
+        ("mlp.shared_experts.", "mlp.shared_mlp."),
+    ]
+    _FUSED_EXPERT_KEY_RE = re.compile(
+        r"^(.*\.mlp\.experts\.)(gate_up_proj|down_proj)$"
+    )
+
+    def _apply_save_reverse_rename_patch():
+        try:
+            from transformers.models.hy_v3.modeling_hy_v3 import HYV3ForCausalLM
+        except ImportError:
+            try:
+                from transformers.hy_v3.modeling_hy_v3 import HYV3ForCausalLM
+            except ImportError:
+                print("[HYV3 Patch] Could not import HYV3ForCausalLM; "
+                      "save reverse rename patch NOT applied.")
+                return
+
+        _orig_save_pretrained = HYV3ForCausalLM.save_pretrained
+
+        def _patched_save_pretrained(self, *args, **kwargs):
+            state_dict = kwargs.get("state_dict", None)
+            if state_dict is not None:
+                reversed_sd = {}
+                for k, v in state_dict.items():
+                    new_k = k
+                    # Apply simple key renames
+                    for new_sub, old_sub in _SAVE_KEY_RENAMES:
+                        if new_sub in new_k:
+                            new_k = new_k.replace(new_sub, old_sub)
+                            break
+
+                    # Check if this is a fused 3D expert key
+                    m = _FUSED_EXPERT_KEY_RE.match(new_k)
+                    if m:
+                        prefix = m.group(1)  # e.g. "model.layers.1.mlp.experts."
+                        proj_type = m.group(2)  # "gate_up_proj" or "down_proj"
+
+                        if proj_type == "gate_up_proj":
+                            # v shape: [num_experts, 2*intermediate, hidden]
+                            num_experts = v.shape[0]
+                            intermediate = v.shape[1] // 2
+                            for i in range(num_experts):
+                                gate = v[i, :intermediate, :]
+                                up = v[i, intermediate:, :]
+                                reversed_sd[f"{prefix}{i}.gate_proj.weight"] = gate
+                                reversed_sd[f"{prefix}{i}.up_proj.weight"] = up
+                        elif proj_type == "down_proj":
+                            # v shape: [num_experts, hidden, intermediate]
+                            num_experts = v.shape[0]
+                            for i in range(num_experts):
+                                reversed_sd[f"{prefix}{i}.down_proj.weight"] = v[i]
+                    else:
+                        reversed_sd[new_k] = v
+
+                kwargs["state_dict"] = reversed_sd
+                print(f"[HYV3 Patch] Reverse-renamed and unfused "
+                      f"{len(state_dict)} -> {len(reversed_sd)} "
+                      f"state_dict keys for old checkpoint compatibility.")
+            return _orig_save_pretrained(self, *args, **kwargs)
+
+        HYV3ForCausalLM.save_pretrained = _patched_save_pretrained
+        print("[HYV3 Patch] Applied: save-time reverse key rename + "
+              "3D -> per-expert unfuse for old ckpt compatibility.")
+
+    _apply_save_reverse_rename_patch()
+    # -------------------------------------------------------------------
+
+    if _has_weights:
+        print(f"Initializing model from local file: {training_args.model_name_or_path}")
+        model = transformers.AutoModelForCausalLM.from_pretrained(
+            training_args.model_name_or_path,
+            trust_remote_code=True,
+            **init_kwargs
+        )
+    else:
+        from transformers import HYV3Config
+        from transformers import HYV3ForCausalLM
+        print(f"Model weights not found at: {training_args.model_name_or_path}, "
+              f"using random initialized HYV3 model instead.")
+        # Use len(tokenizer) to include added special tokens; tokenizer.vocab_size
+        # may only return the base vocabulary size and miss special tokens whose
+        # IDs exceed that range, causing index-out-of-bounds in the embedding layer.
+        config = HYV3Config(
+            vocab_size=len(tokenizer),
+            hidden_size=model_args.hidden_size,
+            intermediate_size=model_args.intermediate_size,
+            max_position_embeddings=training_args.model_max_length,
+            moe_topk=model_args.moe_topk,
+            num_experts=model_args.num_experts,
+            num_attention_heads=model_args.num_attention_heads,
+            num_key_value_heads=model_args.num_key_value_heads,
+            num_hidden_layers=model_args.num_layers,
+            moe_intermediate_size=model_args.moe_intermediate_size,
+            use_mixed_mlp_moe=model_args.use_mixed_mlp_moe,
+            num_shared_expert=model_args.num_shared_expert,
+            use_qk_norm=model_args.use_qk_norm,
+            moe_layer_num_skipped=model_args.moe_layer_num_skipped,
+            tie_word_embeddings=model_args.tie_word_embeddings,
+        )
+        with deepspeed.zero.Init(dtype=init_kwargs.get("torch_dtype", torch.bfloat16), config_dict_or_path=training_args.deepspeed):
+            model = HYV3ForCausalLM(config)
+    
+    if model_args.train_attention_params_only:
+        for name, param in model.named_parameters():
+            if 'self_attn' not in name:
+                param.requires_grad = False
+
+    if model_args.use_lora:
+        # define Lora configuration
+        lora_config = LoraConfig(
+            r=model_args.lora_rank,
+            lora_alpha=model_args.lora_alpha,
+            lora_dropout=model_args.lora_dropout,
+            target_modules=["q_proj", "k_proj", "v_proj", "o_proj"],
+            bias="none",
+            task_type="CAUSAL_LM",
+        )
+        model = get_peft_model(model, lora_config)
+
+    data_module = make_supervised_data_module(tokenizer=tokenizer, data_args=data_args)
+    # Tell Trainer not to attempt DataParallel
+    model.is_parallelizable = True
+    model.model_parallel = True
+
+    training_args.lr_scheduler_kwargs = {
+        'min_lr_rate': training_args.min_lr / training_args.learning_rate,
+    }
+
+    # -----------------------------------------------------------------------
+    # Fix: DeepSpeed ZeRO-3 + gradient checkpointing compatibility.
+    #
+    # PyTorch's torch.utils.checkpoint with use_reentrant=False (the default
+    # in transformers) performs strict metadata checks on recomputed tensors
+    # during backward.  Under ZeRO-3, parameters are all-gathered during the
+    # first forward pass (shape=[full_size]) but may be partitioned back
+    # (shape=[0]) when the checkpoint recomputes, causing a CheckpointError.
+    #
+    # Setting use_reentrant=True avoids this strict metadata check.
+    # -----------------------------------------------------------------------
+    if training_args.gradient_checkpointing and training_args.deepspeed:
+        training_args.gradient_checkpointing_kwargs = {"use_reentrant": True}
+
+    trainer = Trainer(
+        model=model, 
+        processing_class=tokenizer, 
+        args=training_args,
+        callbacks=[CustomSaveCallback],
+        **data_module
+    )
+    model.config.use_cache = False
+
+    # -----------------------------------------------------------------------
+    # Monkey-patch: fix dtype mismatch in DeepSpeed ZeRO-3 linear wrapper.
+    #
+    # By this point the DeepSpeed engine has been initialised by the Trainer
+    # and torch.nn.functional.linear has been replaced with
+    # zero3_linear_wrap.  That wrapper does NOT auto-align input/weight
+    # dtypes before the matmul, causing "expected mat1 and mat2 to have the
+    # same dtype" errors in mixed-precision paths (MoE router gate in fp32
+    # with bf16 weights, expert FFN receiving fp32 routing-weighted input
+    # with bf16 weights, etc.).
+    #
+    # We wrap F.linear HERE (after DeepSpeed init) so that:
+    #   1. We are sure to capture the already-replaced function.
+    #   2. The dtype cast happens *outside* the autograd.Function, so
+    #      gradient-checkpointing recompute sees identical tensor metadata.
+    # -----------------------------------------------------------------------
+    import torch.nn.functional as _F
+    _orig_F_linear = _F.linear
+
+    def _dtype_safe_linear(input, weight, bias=None):
+        if input.dtype != weight.dtype:
+            input = input.to(weight.dtype)
+        return _orig_F_linear(input, weight, bias)
+
+    _F.linear = _dtype_safe_linear
+    # -----------------------------------------------------------------------
+
+    trainer.train(resume_from_checkpoint=training_args.resume_from_checkpoint)
+
+
+if __name__ == "__main__":
+    train()

train/deepspeed_support/train.sh

@@ -0,0 +1,120 @@
+#!/bin/bash
+
+NET_TYPE="high"
+export NCCL_DEBUG=WARN
+export NCCL_P2P_LEVEL=NVL
+export NCCL_IB_TIMEOUT=24
+export NCCL_NVLS_ENABLE=0
+export NCCL_MPI_PROFILE_PRIMS_ENABLE=0
+export CUDA_DEVICE_MAX_CONNECTIONS=1
+export TORCH_NCCL_HEARTBEAT_TIMEOUT_SEC=3600
+if [[ "${NET_TYPE}" = "low" ]]; then
+    export NCCL_SOCKET_IFNAME=eth1
+    export NCCL_IB_GID_INDEX=3
+    export NCCL_IB_HCA=mlx5_2:1
+    export NCCL_IB_SL=3
+    export NCCL_CHECK_DISABLE=1
+    export NCCL_P2P_DISABLE=0
+    export NCCL_LL_THRESHOLD=16384
+    export NCCL_IB_CUDA_SUPPORT=1
+else
+    export NCCL_IB_GID_INDEX=3
+    export NCCL_IB_SL=3
+    export NCCL_CHECK_DISABLE=1
+    export NCCL_P2P_DISABLE=0
+    export NCCL_IB_DISABLE=0
+    export NCCL_LL_THRESHOLD=16384
+    export NCCL_IB_CUDA_SUPPORT=1
+    export NCCL_SOCKET_IFNAME=bond1
+    export UCX_NET_DEVICES=bond1
+    export NCCL_IB_HCA=mlx5_bond_1,mlx5_bond_5,mlx5_bond_3,mlx5_bond_7,mlx5_bond_4,mlx5_bond_8,mlx5_bond_2,mlx5_bond_6
+    export NCCL_COLLNET_ENABLE=0
+    export SHARP_COLL_ENABLE_SAT=0
+    export NCCL_NET_GDR_LEVEL=2
+    export NCCL_IB_QPS_PER_CONNECTION=4
+    export NCCL_IB_TC=160
+    export NCCL_PXN_DISABLE=1
+fi
+
+export HOST_GPU_NUM=8
+# IP list, comma separated. e.g. "192.168.1.1,192.168.1.2" or single node "192.168.1.1"
+IP_LIST=${IP_LIST:-"127.0.0.1"}
+
+IFS=',' read -ra IP_ARRAY <<< "$IP_LIST"
+export NODES=${#IP_ARRAY[@]}
+export LOCAL_IP=${IP_ARRAY[0]}
+NODE_IP_LIST=""
+for ip in "${IP_ARRAY[@]}"; do
+    if [ -n "$NODE_IP_LIST" ]; then
+        NODE_IP_LIST="${NODE_IP_LIST},"
+    fi
+    NODE_IP_LIST="${NODE_IP_LIST}${ip}:${HOST_GPU_NUM}"
+done
+export NODE_IP_LIST
+export NODE_NUM=$((${NODES} * ${HOST_GPU_NUM}))
+
+SCRIPT_DIR=$(dirname "$0")
+
+model_path=path_to_model_weight
+tokenizer_path=../models
+train_data_file=example_data.jsonl
+
+# ds_config_file=ds_zero2_no_offload.json
+# ds_config_file=ds_zero3_no_offload.json
+ds_config_file=${SCRIPT_DIR}/ds_zero3_offload_no_auto.json
+
+output_path=/root/hf_train_output_full
+
+mkdir -p ${output_path}
+
+current_time=$(date "+%Y.%m.%d-%H.%M.%S")
+log_file=${output_path}/"log_${current_time}.txt"
+
+echo $NODE_IP_LIST > env.txt 2>&1
+sed "s/:/ slots=/g" env.txt | sed "s/,/\n/g" >  "hostfile"
+sed "s/:.//g" env.txt | sed "s/,/\n/g" >  "pssh.hosts"
+export CHIEF_IP=$LOCAL_IP
+
+if [ ${NODES} -gt 1 ]; then
+    HOST_PATH=hostfile
+    DS_ARGS="--hostfile=${HOST_PATH} --master_addr ${CHIEF_IP}"
+else
+    DS_ARGS=""
+fi
+
+echo "NODES: ${NODES}, LOCAL_IP: ${LOCAL_IP}, NODE_IP_LIST: ${NODE_IP_LIST}"
+
+deepspeed ${DS_ARGS} \
+    ${SCRIPT_DIR}/train.py \
+    --do_train \
+    --model_name_or_path ${model_path} \
+    --tokenizer_name_or_path ${tokenizer_path} \
+    --train_data_file ${train_data_file} \
+    --deepspeed ${ds_config_file} \
+    --output_dir ${output_path} \
+    --per_device_train_batch_size 1 \
+    --gradient_accumulation_steps 1 \
+    --gradient_checkpointing \
+    --lr_scheduler_type cosine_with_min_lr \
+    --logging_steps 1 \
+    --max_steps 50 \
+    --save_steps 50 \
+    --learning_rate 1e-5 \
+    --min_lr 1e-6 \
+    --warmup_ratio 0.01 \
+    --save_strategy steps \
+    --bf16 \
+    --hidden_size 2048 \
+    --intermediate_size 6912 \
+    --model_max_length 262144 \
+    --max_seq_length 8192 \
+    --moe_topk 8 \
+    --num_experts 128 \
+    --moe_intermediate_size 768 \
+    --moe_layer_num_skipped 1 \
+    --num_attention_heads 32 \
+    --num_key_value_heads 4 \
+    --num_layers 48 \
+    --use_mixed_mlp_moe \
+    --num_shared_expert 1 \
+    --use_qk_norm | tee ${log_file}

train/deepspeed_support/train_dense.py

@@ -0,0 +1,360 @@
+# Copyright 2024 Tencent Inc. 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.
+
+# Copyright 2022 EleutherAI and the HuggingFace Inc. team. All rights reserved.
+#
+# This code is based on EleutherAI's GPT-NeoX library and the GPT-NeoX
+# and OPT implementations in this library. It has been modified from its
+# original forms to accommodate minor architectural differences compared
+# to GPT-NeoX and OPT used by the Meta AI team that trained the model.
+#
+# 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.
+
+"""
+Training script for HunYuan Dense models (1.8B, 7B).
+
+This script is adapted from the original finetune.py for dense models,
+with improvements from the new training framework (train.py for MoE models).
+
+Key differences from train.py (MoE version):
+  - No MoE-related patches (router dtype fix, expert key rename, etc.)
+  - Supports model_size parameter to handle different tokenizer formats
+  - 7B model uses different special tokens than 1.8B model
+"""
+
+import os
+import sys
+sys.path.insert(0, os.path.dirname(os.path.dirname(os.path.abspath(__file__))))
+import json
+import torch
+import shutil
+import logging
+from dataclasses import dataclass, field
+from typing import Optional, Dict, Literal
+
+import transformers
+from torch.utils.data import Dataset
+from transformers import Trainer, TrainerCallback
+from peft import LoraConfig, get_peft_model, PeftModel
+from transformers.trainer_utils import PREFIX_CHECKPOINT_DIR
+from transformers.modeling_utils import unwrap_model
+
+
+def print_args(args, name='arguments'):
+    """Print arguments."""
+    if torch.distributed.get_rank() == 0:
+        print(f'------------------------ {name} ------------------------', flush=True)
+        str_list = []
+        for arg in vars(args):
+            dots = '.' * (48 - len(arg))
+            str_list.append('  {} {} {}'.format(arg, dots, getattr(args, arg)))
+        for arg in sorted(str_list, key=lambda x: x.lower()):
+            print(arg, flush=True)
+        print(f'-------------------- end of {name} ---------------------', flush=True)
+
+
+@dataclass
+class ModelArguments:
+    use_flash_attn: bool = field(
+        default=False,
+        metadata={"help": "Enable FlashAttention-2 for faster training."}
+    )
+    use_lora: bool = field(default=False, metadata={"help": "Enable Lora for faster training."})
+    hidden_size: int = field(default=2048, metadata={"help": "The hidden size of the model."})
+    num_layers: int = field(default=32, metadata={"help": "The number of layers of the model."})
+    num_attention_heads: int = field(default=16, metadata={"help": "The number of attention heads of the model."})
+    intermediate_size: int = field(default=6144, metadata={"help": "The intermediate size of the model."})
+    num_key_value_heads: int = field(default=4, metadata={"help": "The number of key-value heads in GQA."})
+    use_qk_norm: bool = field(default=False, metadata={"help": "Whether to use qk norm."})
+    tie_word_embeddings: bool = field(
+        default=True,
+        metadata={"help": "Whether to tie the word embeddings of the encoder and the decoder."}
+    )
+    lora_rank: int = field(default=64, metadata={"help": "The rank of lora."})
+    lora_alpha: int = field(default=128, metadata={"help": "Lora alpha"})
+    lora_dropout: float = field(default=0.0, metadata={"help": "Lora dropout"})
+    train_attention_params_only: bool = field(default=False, metadata={
+        "help": "Whether to train attention parameters only."}
+    )
+
+
+@dataclass
+class DataArguments:
+    model_size: Literal["0.5B", "1.8B", "4B", "7B"] = field(
+        default="1.8B",
+        metadata={"help": "Select the model size from ['0.5B', '1.8B', '4B', '7B']. "
+                  "This affects the tokenizer special tokens used for loss masking."}
+    )
+    train_data_file: str = field(default=None, metadata={"help": "Path to the training data."})
+    max_seq_length: int = field(
+        default=4096,
+        metadata={"help": "The max sequence length of the model inputs after tokenization."}
+    )
+    use_dummy_data: bool = field(default=False, metadata={"help": "Use dummy data."})
+
+
+@dataclass
+class TrainingArguments(transformers.TrainingArguments):
+    cache_dir: Optional[str] = field(default=None)
+    optim: str = field(default="adamw_torch")
+    model_max_length: int = field(
+        default=4096,
+        metadata={"help": "Maximum sequence length. Sequences will be right padded (and possibly truncated)."},
+    )
+    tokenizer_name_or_path: Optional[str] = field(default=None)
+    model_name_or_path: Optional[str] = field(default=None)
+    min_lr: float = field(
+        default=1e-6,
+        metadata={"help": "The minimum learning rate at the end of the cosine decay."}
+    )
+
+
+IGNORE_INDEX = -100
+
+
+class DummyDataset(Dataset):
+    def __init__(self, tokenizer, max_seq_length=512, length=1000):
+        self.tokenizer = tokenizer
+        self.max_seq_length = max_seq_length
+        self.length = length
+
+    def __len__(self):
+        return self.length
+
+    def __getitem__(self, index):
+        tokens = torch.randint(0, self.tokenizer.vocab_size, (self.max_seq_length,))
+        return {'input_ids': tokens, 'labels': tokens}
+
+
+class SFTDataset(Dataset):
+    def __init__(self, data_file, tokenizer, max_seq_length=4096, model_size="1.8B"):
+        self.tokenizer = tokenizer
+        self.max_seq_length = max_seq_length
+        self.model_size = model_size
+        self.data_list = self.load_data(data_file)
+
+    def __len__(self):
+        return len(self.data_list)
+
+    def load_data(self, data_file):
+        logging.info('Loading data: {}'.format(data_file))
+        with open(data_file, 'r', encoding='utf8') as f:
+            data_list = f.readlines()
+        logging.info("there are {} data in dataset".format(len(data_list)))
+        return data_list
+
+    def encode_data(self, data_dict):
+        model_inputs = {}
+        template_output = self.tokenizer.apply_chat_template(
+            data_dict['messages'], tokenize=True, return_dict=False
+        )
+        if isinstance(template_output, list) and len(template_output) > 0 and isinstance(template_output[0], list):
+            template_output = template_output[0]
+        message_tokens = torch.tensor(template_output, dtype=torch.long)
+
+        # Note: The 7B model uses a different vocabulary/special tokens than other models.
+        if self.model_size == "7B":
+            sep_token_id = self.tokenizer.convert_tokens_to_ids('<|extra_0|>')
+            eos_token_id = self.tokenizer.convert_tokens_to_ids('<|eos|>')
+        else:
+            sep_token_id = self.tokenizer.convert_tokens_to_ids('<｜hy_Assistant｜>')
+            eos_token_id = self.tokenizer.convert_tokens_to_ids('<｜hy_place▁holder▁no▁2｜>')
+
+        # Find assistant reply boundaries
+        loss_token_begins = (message_tokens == sep_token_id).nonzero(as_tuple=True)[0].tolist()
+        loss_token_ends = (message_tokens == eos_token_id).nonzero(as_tuple=True)[0].tolist()
+        message_labels = torch.tensor([IGNORE_INDEX] * message_tokens.shape[0])
+        for begin_idx, end_idx in zip(loss_token_begins, loss_token_ends):
+            # Compute loss from sep_token to eos_token (inclusive)
+            message_labels[begin_idx:end_idx + 1] = message_tokens[begin_idx:end_idx + 1]
+
+        input_ids = message_tokens.to(torch.long)
+        labels = message_labels.to(torch.long)
+
+        input_ids = input_ids[:self.max_seq_length]
+        labels = labels[:self.max_seq_length]
+
+        pad_token_id = self.tokenizer.pad_token_id
+        attention_mask = [1 if val != pad_token_id else 0 for val in input_ids]
+        model_inputs["input_ids"] = input_ids
+        model_inputs["attention_mask"] = torch.tensor(attention_mask, dtype=torch.bool)
+        model_inputs["labels"] = labels
+
+        return model_inputs
+
+    def __getitem__(self, index):
+        data = self.data_list[index]
+        data = json.loads(data)
+        model_inputs = self.encode_data(data)
+        return model_inputs
+
+
+@dataclass
+class DataCollatorForSupervisedDataset(object):
+    """Collate examples for supervised fine-tuning."""
+
+    tokenizer: transformers.PreTrainedTokenizer
+
+    def __call__(self, instances):
+        input_ids = [instance['input_ids'] for instance in instances]
+        labels = [instance['labels'] for instance in instances]
+        pad_token_id = self.tokenizer.pad_token_id
+        input_ids = torch.nn.utils.rnn.pad_sequence(input_ids, batch_first=True, padding_value=pad_token_id)
+        labels = torch.nn.utils.rnn.pad_sequence(labels, batch_first=True, padding_value=IGNORE_INDEX)
+        return dict(
+            input_ids=input_ids,
+            labels=labels,
+            attention_mask=input_ids.ne(pad_token_id),
+        )
+
+
+def make_supervised_data_module(tokenizer, data_args) -> Dict:
+    """Make dataset and collator for supervised fine-tuning."""
+    if data_args.use_dummy_data:
+        train_dataset = DummyDataset(tokenizer, data_args.max_seq_length)
+    else:
+        train_dataset = SFTDataset(
+            tokenizer=tokenizer,
+            data_file=data_args.train_data_file,
+            max_seq_length=data_args.max_seq_length,
+            model_size=data_args.model_size,
+        )
+    data_collator = DataCollatorForSupervisedDataset(tokenizer=tokenizer)
+    return dict(train_dataset=train_dataset, eval_dataset=None, data_collator=data_collator)
+
+
+# Copy tokenizer and config files to each checkpoint directory for self-contained inference
+class CustomSaveCallback(TrainerCallback):
+    def on_save(self, args, state, control, **kwargs):
+        if torch.distributed.get_rank() == 0:
+            output_dir = os.path.join(args.output_dir, f"{PREFIX_CHECKPOINT_DIR}-{state.global_step}")
+
+            # Copy tokenizer files to checkpoint directory
+            tokenizer_files = [
+                'generation_config.json',
+                'hy.tiktoken',
+                'tokenizer_config.json',
+                'tokenization_hy.py',
+                'tokenizer.json',
+                'special_tokens_map.json',
+                'chat_template.jinja',
+                'config.json',
+            ]
+            src_dir = args.tokenizer_name_or_path or args.model_name_or_path
+            for fname in tokenizer_files:
+                src = os.path.join(src_dir, fname)
+                if os.path.isfile(src):
+                    shutil.copy(src, os.path.join(output_dir, fname))
+
+        return control
+
+
+def train():
+    parser = transformers.HfArgumentParser((ModelArguments, DataArguments, TrainingArguments))
+    model_args, data_args, training_args = parser.parse_args_into_dataclasses()
+    print_args(model_args, 'model arguments')
+    print_args(data_args, 'data arguments')
+    print_args(training_args, 'training arguments')
+
+    tokenizer = transformers.AutoTokenizer.from_pretrained(
+        training_args.tokenizer_name_or_path,
+        trust_remote_code=True
+    )
+
+    init_kwargs = {}
+    if model_args.use_flash_attn:
+        init_kwargs["attn_implementation"] = "flash_attention_2"
+    if training_args.bf16:
+        init_kwargs["torch_dtype"] = torch.bfloat16
+    elif training_args.fp16:
+        init_kwargs["torch_dtype"] = torch.float16
+
+    # Load model from pretrained weights
+    if training_args.model_name_or_path is not None and os.path.exists(training_args.model_name_or_path):
+        print(f"Initializing model from local file: {training_args.model_name_or_path}")
+        model = transformers.AutoModelForCausalLM.from_pretrained(
+            training_args.model_name_or_path,
+            trust_remote_code=True,
+            **init_kwargs
+        )
+    else:
+        raise FileNotFoundError(
+            f"Model path {training_args.model_name_or_path} is invalid or does not exist. "
+            f"Dense model training requires pre-trained weights."
+        )
+
+    if model_args.train_attention_params_only:
+        for name, param in model.named_parameters():
+            if 'self_attn' not in name:
+                param.requires_grad = False
+
+    if model_args.use_lora:
+        # Define LoRA configuration
+        lora_config = LoraConfig(
+            r=model_args.lora_rank,
+            lora_alpha=model_args.lora_alpha,
+            lora_dropout=model_args.lora_dropout,
+            target_modules=["q_proj", "k_proj", "v_proj", "o_proj"],
+            bias="none",
+            task_type="CAUSAL_LM",
+        )
+        model = get_peft_model(model, lora_config)
+
+    data_module = make_supervised_data_module(tokenizer=tokenizer, data_args=data_args)
+    # Tell Trainer not to attempt DataParallel
+    model.is_parallelizable = True
+    model.model_parallel = True
+
+    training_args.lr_scheduler_kwargs = {
+        'min_lr_rate': training_args.min_lr / training_args.learning_rate,
+    }
+
+    # -----------------------------------------------------------------------
+    # Fix: DeepSpeed ZeRO-3 + gradient checkpointing compatibility.
+    #
+    # PyTorch's torch.utils.checkpoint with use_reentrant=False (the default
+    # in transformers) performs strict metadata checks on recomputed tensors
+    # during backward.  Under ZeRO-3, parameters are all-gathered during the
+    # first forward pass (shape=[full_size]) but may be partitioned back
+    # (shape=[0]) when the checkpoint recomputes, causing a CheckpointError.
+    #
+    # Setting use_reentrant=True avoids this strict metadata check.
+    # -----------------------------------------------------------------------
+    if training_args.gradient_checkpointing and training_args.deepspeed:
+        training_args.gradient_checkpointing_kwargs = {"use_reentrant": True}
+
+    trainer = Trainer(
+        model=model,
+        processing_class=tokenizer,
+        args=training_args,
+        callbacks=[CustomSaveCallback],
+        **data_module
+    )
+    model.config.use_cache = False
+
+    trainer.train(resume_from_checkpoint=training_args.resume_from_checkpoint)
+
+
+if __name__ == "__main__":
+    train()

train/deepspeed_support/train_dense.sh

@@ -0,0 +1,155 @@
+#!/bin/bash
+
+# Unified Dense model full fine-tuning script
+# Supports: 1.8B and 7B dense models
+# Usage: bash train_dense.sh [1.8B|7B]
+#   - 1.8B: 1x GPU (24GB+), DeepSpeed ZeRO-2 (no offload)
+#   - 7B:   2x GPU (80GB+ each), DeepSpeed ZeRO-3 (no offload)
+
+# ============== Model Size Selection ==============
+MODEL_SIZE=${1:-"1.8B"}
+
+if [[ "${MODEL_SIZE}" != "1.8B" && "${MODEL_SIZE}" != "7B" ]]; then
+    echo "Error: MODEL_SIZE must be '1.8B' or '7B', got '${MODEL_SIZE}'"
+    echo "Usage: bash train_dense.sh [1.8B|7B]"
+    exit 1
+fi
+
+# ============== NCCL Configuration ==============
+NET_TYPE="high"
+export NCCL_DEBUG=WARN
+export NCCL_P2P_LEVEL=NVL
+export NCCL_IB_TIMEOUT=24
+export NCCL_NVLS_ENABLE=0
+export NCCL_MPI_PROFILE_PRIMS_ENABLE=0
+export CUDA_DEVICE_MAX_CONNECTIONS=1
+export TORCH_NCCL_HEARTBEAT_TIMEOUT_SEC=3600
+if [[ "${NET_TYPE}" = "low" ]]; then
+    export NCCL_SOCKET_IFNAME=eth1
+    export NCCL_IB_GID_INDEX=3
+    export NCCL_IB_HCA=mlx5_2:1
+    export NCCL_IB_SL=3
+    export NCCL_CHECK_DISABLE=1
+    export NCCL_P2P_DISABLE=0
+    export NCCL_LL_THRESHOLD=16384
+    export NCCL_IB_CUDA_SUPPORT=1
+else
+    export NCCL_IB_GID_INDEX=3
+    export NCCL_IB_SL=3
+    export NCCL_CHECK_DISABLE=1
+    export NCCL_P2P_DISABLE=0
+    export NCCL_IB_DISABLE=0
+    export NCCL_LL_THRESHOLD=16384
+    export NCCL_IB_CUDA_SUPPORT=1
+    export NCCL_SOCKET_IFNAME=bond1
+    export UCX_NET_DEVICES=bond1
+    export NCCL_IB_HCA=mlx5_bond_1,mlx5_bond_5,mlx5_bond_3,mlx5_bond_7,mlx5_bond_4,mlx5_bond_8,mlx5_bond_2,mlx5_bond_6
+    export NCCL_COLLNET_ENABLE=0
+    export SHARP_COLL_ENABLE_SAT=0
+    export NCCL_NET_GDR_LEVEL=2
+    export NCCL_IB_QPS_PER_CONNECTION=4
+    export NCCL_IB_TC=160
+    export NCCL_PXN_DISABLE=1
+fi
+
+# ============== Model-specific Configuration ==============
+SCRIPT_DIR=$(dirname "$0")
+
+if [[ "${MODEL_SIZE}" == "1.8B" ]]; then
+    export HOST_GPU_NUM=1
+    model_path=path_to_dense_1_8b_model
+    ds_config_file=${SCRIPT_DIR}/ds_zero2_no_offload.json
+    output_path=./dense_1_8b_output
+    HIDDEN_SIZE=2048
+    INTERMEDIATE_SIZE=6144
+    NUM_ATTENTION_HEADS=16
+    NUM_KEY_VALUE_HEADS=4
+    NUM_LAYERS=32
+else
+    export HOST_GPU_NUM=2
+    model_path=path_to_dense_7b_model
+    ds_config_file=${SCRIPT_DIR}/ds_zero3_no_offload.json
+    output_path=./dense_7b_output
+    HIDDEN_SIZE=4096
+    INTERMEDIATE_SIZE=14336
+    NUM_ATTENTION_HEADS=32
+    NUM_KEY_VALUE_HEADS=8
+    NUM_LAYERS=32
+fi
+
+tokenizer_path=${model_path}
+train_data_file=../data/example_data.jsonl
+
+# ============== Multi-node Configuration ==============
+# IP list, comma separated. e.g. "192.168.1.1,192.168.1.2" or single node "192.168.1.1"
+IP_LIST=${IP_LIST:-"127.0.0.1"}
+
+IFS=',' read -ra IP_ARRAY <<< "$IP_LIST"
+export NODES=${#IP_ARRAY[@]}
+export LOCAL_IP=${IP_ARRAY[0]}
+NODE_IP_LIST=""
+for ip in "${IP_ARRAY[@]}"; do
+    if [ -n "$NODE_IP_LIST" ]; then
+        NODE_IP_LIST="${NODE_IP_LIST},"
+    fi
+    NODE_IP_LIST="${NODE_IP_LIST}${ip}:${HOST_GPU_NUM}"
+done
+export NODE_IP_LIST
+export NODE_NUM=$((${NODES} * ${HOST_GPU_NUM}))
+
+# ============== Output & Logging ==============
+mkdir -p ${output_path}
+
+current_time=$(date "+%Y.%m.%d-%H.%M.%S")
+log_file=${output_path}/"log_${current_time}.txt"
+
+echo $NODE_IP_LIST > env.txt 2>&1
+sed "s/:/ slots=/g" env.txt | sed "s/,/\n/g" >  "hostfile"
+sed "s/:.//g" env.txt | sed "s/,/\n/g" >  "pssh.hosts"
+export CHIEF_IP=$LOCAL_IP
+
+if [ ${NODES} -gt 1 ]; then
+    HOST_PATH=hostfile
+    DS_ARGS="--hostfile=${HOST_PATH} --master_addr ${CHIEF_IP}"
+else
+    DS_ARGS=""
+fi
+
+echo "============================================"
+echo "Dense ${MODEL_SIZE} full fine-tuning"
+echo "NODES: ${NODES}, LOCAL_IP: ${LOCAL_IP}, NODE_IP_LIST: ${NODE_IP_LIST}"
+echo "DeepSpeed config: ${ds_config_file}"
+echo "Model path: ${model_path}"
+echo "Output path: ${output_path}"
+echo "============================================"
+
+# ============== Launch Training ==============
+deepspeed ${DS_ARGS} \
+    ${SCRIPT_DIR}/train_dense.py \
+    --do_train \
+    --model_size ${MODEL_SIZE} \
+    --model_name_or_path ${model_path} \
+    --tokenizer_name_or_path ${tokenizer_path} \
+    --train_data_file ${train_data_file} \
+    --deepspeed ${ds_config_file} \
+    --output_dir ${output_path} \
+    --per_device_train_batch_size 1 \
+    --gradient_accumulation_steps 1 \
+    --gradient_checkpointing \
+    --lr_scheduler_type cosine_with_min_lr \
+    --logging_steps 1 \
+    --max_steps 30 \
+    --save_steps 30 \
+    --learning_rate 1e-5 \
+    --min_lr 1e-6 \
+    --warmup_ratio 0.01 \
+    --save_strategy steps \
+    --bf16 \
+    --hidden_size ${HIDDEN_SIZE} \
+    --intermediate_size ${INTERMEDIATE_SIZE} \
+    --num_attention_heads ${NUM_ATTENTION_HEADS} \
+    --num_key_value_heads ${NUM_KEY_VALUE_HEADS} \
+    --num_layers ${NUM_LAYERS} \
+    --model_max_length 4096 \
+    --max_seq_length 4096 \
+    --use_qk_norm | tee ${log_file}

train/deepspeed_support/train_dense_lora.sh

@@ -0,0 +1,161 @@
+#!/bin/bash
+
+# Unified Dense model LoRA fine-tuning script
+# Supports: 1.8B and 7B dense models
+# Usage: bash train_dense_lora.sh [1.8B|7B]
+#   - 1.8B: 1x GPU (24GB+), DeepSpeed ZeRO-2 (no offload)
+#   - 7B:   1x GPU (80GB+), DeepSpeed ZeRO-2 (no offload)
+# LoRA greatly reduces memory requirements compared to full fine-tuning.
+
+# ============== Model Size Selection ==============
+MODEL_SIZE=${1:-"1.8B"}
+
+if [[ "${MODEL_SIZE}" != "1.8B" && "${MODEL_SIZE}" != "7B" ]]; then
+    echo "Error: MODEL_SIZE must be '1.8B' or '7B', got '${MODEL_SIZE}'"
+    echo "Usage: bash train_dense_lora.sh [1.8B|7B]"
+    exit 1
+fi
+
+# ============== NCCL Configuration ==============
+NET_TYPE="high"
+export NCCL_DEBUG=WARN
+export NCCL_P2P_LEVEL=NVL
+export NCCL_IB_TIMEOUT=24
+export NCCL_NVLS_ENABLE=0
+export NCCL_MPI_PROFILE_PRIMS_ENABLE=0
+export CUDA_DEVICE_MAX_CONNECTIONS=1
+export TORCH_NCCL_HEARTBEAT_TIMEOUT_SEC=3600
+if [[ "${NET_TYPE}" = "low" ]]; then
+    export NCCL_SOCKET_IFNAME=eth1
+    export NCCL_IB_GID_INDEX=3
+    export NCCL_IB_HCA=mlx5_2:1
+    export NCCL_IB_SL=3
+    export NCCL_CHECK_DISABLE=1
+    export NCCL_P2P_DISABLE=0
+    export NCCL_LL_THRESHOLD=16384
+    export NCCL_IB_CUDA_SUPPORT=1
+else
+    export NCCL_IB_GID_INDEX=3
+    export NCCL_IB_SL=3
+    export NCCL_CHECK_DISABLE=1
+    export NCCL_P2P_DISABLE=0
+    export NCCL_IB_DISABLE=0
+    export NCCL_LL_THRESHOLD=16384
+    export NCCL_IB_CUDA_SUPPORT=1
+    export NCCL_SOCKET_IFNAME=bond1
+    export UCX_NET_DEVICES=bond1
+    export NCCL_IB_HCA=mlx5_bond_1,mlx5_bond_5,mlx5_bond_3,mlx5_bond_7,mlx5_bond_4,mlx5_bond_8,mlx5_bond_2,mlx5_bond_6
+    export NCCL_COLLNET_ENABLE=0
+    export SHARP_COLL_ENABLE_SAT=0
+    export NCCL_NET_GDR_LEVEL=2
+    export NCCL_IB_QPS_PER_CONNECTION=4
+    export NCCL_IB_TC=160
+    export NCCL_PXN_DISABLE=1
+fi
+
+# ============== Model-specific Configuration ==============
+SCRIPT_DIR=$(dirname "$0")
+
+# LoRA training uses ZeRO-2 (no offload) for both 1.8B and 7B
+# since only adapter parameters are trained, memory usage is much lower
+export HOST_GPU_NUM=1
+ds_config_file=${SCRIPT_DIR}/ds_zero2_no_offload.json
+
+if [[ "${MODEL_SIZE}" == "1.8B" ]]; then
+    model_path=path_to_dense_1_8b_model
+    output_path=dense_1_8b_lora_output
+    HIDDEN_SIZE=2048
+    INTERMEDIATE_SIZE=6144
+    NUM_ATTENTION_HEADS=16
+    NUM_KEY_VALUE_HEADS=4
+    NUM_LAYERS=32
+else
+    model_path=path_to_dense_7b_model
+    output_path=dense_7b_lora_output
+    HIDDEN_SIZE=4096
+    INTERMEDIATE_SIZE=14336
+    NUM_ATTENTION_HEADS=32
+    NUM_KEY_VALUE_HEADS=8
+    NUM_LAYERS=32
+fi
+
+tokenizer_path=${model_path}
+train_data_file=${SCRIPT_DIR}/../data/example_data.jsonl
+
+# ============== Multi-node Configuration ==============
+# IP list, comma separated. e.g. "192.168.1.1,192.168.1.2" or single node "192.168.1.1"
+IP_LIST=${IP_LIST:-"127.0.0.1"}
+
+IFS=',' read -ra IP_ARRAY <<< "$IP_LIST"
+export NODES=${#IP_ARRAY[@]}
+export LOCAL_IP=${IP_ARRAY[0]}
+NODE_IP_LIST=""
+for ip in "${IP_ARRAY[@]}"; do
+    if [ -n "$NODE_IP_LIST" ]; then
+        NODE_IP_LIST="${NODE_IP_LIST},"
+    fi
+    NODE_IP_LIST="${NODE_IP_LIST}${ip}:${HOST_GPU_NUM}"
+done
+export NODE_IP_LIST
+export NODE_NUM=$((${NODES} * ${HOST_GPU_NUM}))
+
+# ============== Output & Logging ==============
+mkdir -p ${output_path}
+
+current_time=$(date "+%Y.%m.%d-%H.%M.%S")
+log_file=${output_path}/"log_${current_time}.txt"
+
+echo $NODE_IP_LIST > env.txt 2>&1
+sed "s/:/ slots=/g" env.txt | sed "s/,/\n/g" >  "hostfile"
+sed "s/:.//g" env.txt | sed "s/,/\n/g" >  "pssh.hosts"
+export CHIEF_IP=$LOCAL_IP
+
+if [ ${NODES} -gt 1 ]; then
+    HOST_PATH=hostfile
+    DS_ARGS="--hostfile=${HOST_PATH} --master_addr ${CHIEF_IP}"
+else
+    DS_ARGS=""
+fi
+
+echo "============================================"
+echo "Dense ${MODEL_SIZE} LoRA fine-tuning"
+echo "NODES: ${NODES}, LOCAL_IP: ${LOCAL_IP}, NODE_IP_LIST: ${NODE_IP_LIST}"
+echo "DeepSpeed config: ${ds_config_file}"
+echo "Model path: ${model_path}"
+echo "Output path: ${output_path}"
+echo "============================================"
+
+# ============== Launch Training ==============
+deepspeed ${DS_ARGS} \
+    ${SCRIPT_DIR}/train_dense.py \
+    --do_train \
+    --model_size ${MODEL_SIZE} \
+    --model_name_or_path ${model_path} \
+    --tokenizer_name_or_path ${tokenizer_path} \
+    --train_data_file ${train_data_file} \
+    --deepspeed ${ds_config_file} \
+    --output_dir ${output_path} \
+    --per_device_train_batch_size 1 \
+    --gradient_accumulation_steps 1 \
+    --gradient_checkpointing \
+    --lr_scheduler_type cosine_with_min_lr \
+    --logging_steps 1 \
+    --max_steps 30 \
+    --save_steps 30 \
+    --learning_rate 2e-4 \
+    --min_lr 1e-5 \
+    --warmup_ratio 0.01 \
+    --save_strategy steps \
+    --bf16 \
+    --hidden_size ${HIDDEN_SIZE} \
+    --intermediate_size ${INTERMEDIATE_SIZE} \
+    --num_attention_heads ${NUM_ATTENTION_HEADS} \
+    --num_key_value_heads ${NUM_KEY_VALUE_HEADS} \
+    --num_layers ${NUM_LAYERS} \
+    --model_max_length 4096 \
+    --max_seq_length 4096 \
+    --use_qk_norm \
+    --use_lora \
+    --lora_rank 64 \
+    --lora_alpha 128 \
+    --lora_dropout 0.05 | tee ${log_file}

train/deepspeed_support/train_lora.sh

@@ -0,0 +1,125 @@
+#!/bin/bash
+
+NET_TYPE="high"
+export NCCL_DEBUG=WARN
+export NCCL_P2P_LEVEL=NVL
+export NCCL_IB_TIMEOUT=24
+export NCCL_NVLS_ENABLE=0
+export NCCL_MPI_PROFILE_PRIMS_ENABLE=0
+export CUDA_DEVICE_MAX_CONNECTIONS=1
+export TORCH_NCCL_HEARTBEAT_TIMEOUT_SEC=3600
+if [[ "${NET_TYPE}" = "low" ]]; then
+    export NCCL_SOCKET_IFNAME=eth1
+    export NCCL_IB_GID_INDEX=3
+    export NCCL_IB_HCA=mlx5_2:1
+    export NCCL_IB_SL=3
+    export NCCL_CHECK_DISABLE=1
+    export NCCL_P2P_DISABLE=0
+    export NCCL_LL_THRESHOLD=16384
+    export NCCL_IB_CUDA_SUPPORT=1
+else
+    export NCCL_IB_GID_INDEX=3
+    export NCCL_IB_SL=3
+    export NCCL_CHECK_DISABLE=1
+    export NCCL_P2P_DISABLE=0
+    export NCCL_IB_DISABLE=0
+    export NCCL_LL_THRESHOLD=16384
+    export NCCL_IB_CUDA_SUPPORT=1
+    export NCCL_SOCKET_IFNAME=bond1
+    export UCX_NET_DEVICES=bond1
+    export NCCL_IB_HCA=mlx5_bond_1,mlx5_bond_5,mlx5_bond_3,mlx5_bond_7,mlx5_bond_4,mlx5_bond_8,mlx5_bond_2,mlx5_bond_6
+    export NCCL_COLLNET_ENABLE=0
+    export SHARP_COLL_ENABLE_SAT=0
+    export NCCL_NET_GDR_LEVEL=2
+    export NCCL_IB_QPS_PER_CONNECTION=4
+    export NCCL_IB_TC=160
+    export NCCL_PXN_DISABLE=1
+fi
+
+export HOST_GPU_NUM=8
+# IP list, comma separated. e.g. "192.168.1.1,192.168.1.2" or single node "192.168.1.1"
+IP_LIST=${IP_LIST:-"127.0.0.1"}
+
+IFS=',' read -ra IP_ARRAY <<< "$IP_LIST"
+export NODES=${#IP_ARRAY[@]}
+export LOCAL_IP=${IP_ARRAY[0]}
+NODE_IP_LIST=""
+for ip in "${IP_ARRAY[@]}"; do
+    if [ -n "$NODE_IP_LIST" ]; then
+        NODE_IP_LIST="${NODE_IP_LIST},"
+    fi
+    NODE_IP_LIST="${NODE_IP_LIST}${ip}:${HOST_GPU_NUM}"
+done
+export NODE_IP_LIST
+export NODE_NUM=$((${NODES} * ${HOST_GPU_NUM}))
+
+SCRIPT_DIR=$(dirname "$0")
+
+model_path=path_to_model_weight
+tokenizer_path=../models
+train_data_file=example_data.jsonl
+
+# ds_config_file=ds_zero2_no_offload.json
+# ds_config_file=ds_zero3_no_offload.json
+# For LoRA, zero2_offload is recommended to save memory
+ds_config_file=${SCRIPT_DIR}/ds_zero2_no_offload.json
+
+output_path=/root/hf_train_output
+
+mkdir -p ${output_path}
+
+current_time=$(date "+%Y.%m.%d-%H.%M.%S")
+log_file=${output_path}/"log_${current_time}.txt"
+
+echo $NODE_IP_LIST > env.txt 2>&1
+sed "s/:/ slots=/g" env.txt | sed "s/,/\n/g" >  "hostfile"
+sed "s/:.//g" env.txt | sed "s/,/\n/g" >  "pssh.hosts"
+export CHIEF_IP=$LOCAL_IP
+
+if [ ${NODES} -gt 1 ]; then
+    HOST_PATH=hostfile
+    DS_ARGS="--hostfile=${HOST_PATH} --master_addr ${CHIEF_IP}"
+else
+    DS_ARGS=""
+fi
+
+echo "NODES: ${NODES}, LOCAL_IP: ${LOCAL_IP}, NODE_IP_LIST: ${NODE_IP_LIST}"
+
+deepspeed ${DS_ARGS} \
+    ${SCRIPT_DIR}/train.py \
+    --do_train \
+    --model_name_or_path ${model_path} \
+    --tokenizer_name_or_path ${tokenizer_path} \
+    --train_data_file ${train_data_file} \
+    --deepspeed ${ds_config_file} \
+    --output_dir ${output_path} \
+    --per_device_train_batch_size 1 \
+    --gradient_accumulation_steps 1 \
+    --gradient_checkpointing \
+    --lr_scheduler_type cosine_with_min_lr \
+    --logging_steps 1 \
+    --max_steps 200 \
+    --save_steps 100 \
+    --learning_rate 1e-5 \
+    --min_lr 1e-6 \
+    --warmup_ratio 0.01 \
+    --save_strategy steps \
+    --bf16 \
+    --use_lora \
+    --lora_rank 64 \
+    --lora_alpha 128 \
+    --lora_dropout 0.1 \
+    --hidden_size 2048 \
+    --intermediate_size 6912 \
+    --model_max_length 8192 \
+    --max_seq_length 8192 \
+    --moe_topk 8 \
+    --num_experts 128 \
+    --moe_intermediate_size 768 \
+    --moe_layer_num_skipped 1 \
+    --num_attention_heads 32 \
+    --num_key_value_heads 4 \
+    --num_layers 48 \
+    --use_mixed_mlp_moe \
+    --num_shared_expert 1 \
+    --use_qk_norm | tee ${log_file}

train/llama_factory_support/convert_zero_to_hf.sh

@@ -0,0 +1,175 @@
+#!/bin/bash
+# 将 DeepSpeed ZeRO 格式的 checkpoint 转换为 HuggingFace 格式
+# 使用 zero_to_fp32.py 转换权重，然后保存为 HF 格式
+
+set -e  # 遇到错误立即退出
+
+SCRIPT_DIR="$(cd "$(dirname "$0")" && pwd)"
+PROJECT_ROOT="$(dirname "$(dirname "$SCRIPT_DIR")")"
+
+# 激活 conda 环境
+source "$(conda info --base)/etc/profile.d/conda.sh"
+conda activate llama_factory
+
+# 设置环境变量
+export LD_LIBRARY_PATH=$CONDA_PREFIX/lib:$LD_LIBRARY_PATH
+export DISABLE_VERSION_CHECK=1
+export CUDA_VISIBLE_DEVICES=""  # 使用 CPU 进行转换，避免显存不足
+
+CHECKPOINT_DIR="$SCRIPT_DIR/saves/hy_v3/full/sft/checkpoint-39"
+OUTPUT_DIR="$SCRIPT_DIR/saves/hy_v3/full/sft/checkpoint-39/hf_converted"
+TEMP_WEIGHTS_DIR="$OUTPUT_DIR/zero_fp32_output"  # 分片输出目录
+
+echo "=========================================="
+echo "Converting DeepSpeed ZeRO checkpoint to HF format"
+echo "Input:  $CHECKPOINT_DIR"
+echo "Output: $OUTPUT_DIR"
+echo "=========================================="
+
+# 创建输出目录
+mkdir -p "$OUTPUT_DIR"
+rm -rf "$TEMP_WEIGHTS_DIR"
+mkdir -p "$TEMP_WEIGHTS_DIR"
+
+# Step 1: 使用 zero_to_fp32.py 转换权重
+echo ""
+echo "[Step 1/3] Converting weights from ZeRO format to FP32..."
+cd "$CHECKPOINT_DIR"
+python3 zero_to_fp32.py . "$TEMP_WEIGHTS_DIR"
+
+# 检查输出 - zero_to_fp32.py 可能输出单个文件或多个分片
+if [ -d "$TEMP_WEIGHTS_DIR" ] && [ "$(ls -A "$TEMP_WEIGHTS_DIR" 2>/dev/null)" ]; then
+    echo "Weight conversion completed! Output in: $TEMP_WEIGHTS_DIR"
+    echo "Files: $(ls "$TEMP_WEIGHTS_DIR" | wc -l) files"
+else
+    echo "ERROR: Weight conversion failed!"
+    exit 1
+fi
+
+# Step 2: 复制配置文件
+echo ""
+echo "[Step 2/3] Copying config files..."
+BASE_MODEL_DIR="$PROJECT_ROOT/a3b_ckpt"
+
+# 从基座模型复制完整的配置文件（checkpoint 中的 tokenizer_config 不完整）
+cp "$BASE_MODEL_DIR/config.json" "$OUTPUT_DIR/" 2>/dev/null || true
+cp "$BASE_MODEL_DIR/tokenizer_config.json" "$OUTPUT_DIR/" 2>/dev/null || true
+cp "$BASE_MODEL_DIR/tokenizer.json" "$OUTPUT_DIR/" 2>/dev/null || true
+cp "$BASE_MODEL_DIR/special_tokens_map.json" "$OUTPUT_DIR/" 2>/dev/null || true
+cp "$BASE_MODEL_DIR/chat_template.jinja" "$OUTPUT_DIR/" 2>/dev/null || true
+cp "$CHECKPOINT_DIR/generation_config.json" "$OUTPUT_DIR/" 2>/dev/null || true
+
+echo "Config files copied from base model."
+
+# Step 3: 加载权重并保存为 HF 格式
+echo ""
+echo "[Step 3/3] Converting to HuggingFace format..."
+
+cat > /tmp/convert_to_hf.py << 'PYEOF'
+import torch
+import json
+import os
+import sys
+import glob
+sys.path.insert(0, os.environ.get("PROJECT_ROOT", "."))
+
+# 设置目录
+output_dir = os.environ["OUTPUT_DIR"]
+checkpoint_dir = os.environ["CHECKPOINT_DIR"]
+base_model_dir = os.environ["BASE_MODEL_DIR"]
+temp_weights_dir = os.path.join(output_dir, "zero_fp32_output")
+
+# 加载 tokenizer
+from transformers import AutoTokenizer
+print("Loading tokenizer...")
+tokenizer = AutoTokenizer.from_pretrained(base_model_dir, trust_remote_code=True)
+tokenizer.save_pretrained(output_dir)
+
+# 从 base model 加载配置和模型结构
+from transformers import AutoConfig, AutoModelForCausalLM
+print("Loading config...")
+config = AutoConfig.from_pretrained(base_model_dir, trust_remote_code=True)
+
+# 创建模型（从基座模型加载结构和权重，然后用训练后的权重覆盖）
+print("Loading base model...")
+with torch.no_grad():
+    model = AutoModelForCausalLM.from_pretrained(
+        base_model_dir,
+        torch_dtype=torch.bfloat16,
+        trust_remote_code=True,
+    )
+
+# 加载转换后的权重（可能分片）
+print(f"Loading weights from {temp_weights_dir} ...")
+weight_files = sorted(glob.glob(os.path.join(temp_weights_dir, "*.bin")) + 
+                      glob.glob(os.path.join(temp_weights_dir, "*.safetensors")))
+
+if not weight_files:
+    print(f"ERROR: No weight files found in {temp_weights_dir}")
+    sys.exit(1)
+
+print(f"Found {len(weight_files)} weight files")
+state_dict = {}
+for wf in weight_files:
+    print(f"  Loading {wf} ...")
+    if wf.endswith('.safetensors'):
+        from safetensors.torch import load_file
+        state_dict.update(load_file(wf, device="cpu"))
+    else:
+        state_dict.update(torch.load(wf, map_location="cpu"))
+
+# 加载权重到模型
+print("Loading converted weights into model...")
+model_state_dict = model.state_dict()
+filtered_state_dict = {}
+skipped = 0
+matched = 0
+for k, v in state_dict.items():
+    if k in model_state_dict:
+        # 转换 dtype
+        if v.dtype != model_state_dict[k].dtype:
+            v = v.to(model_state_dict[k].dtype)
+        filtered_state_dict[k] = v
+        matched += 1
+    else:
+        skipped += 1
+        if skipped <= 10:  # 只打印前10个跳过的key
+            print(f"  Skipping key: {k}")
+
+print(f"Matched {matched} tensors, skipped {skipped}")
+
+if matched == 0:
+    print("ERROR: No weights matched! Something is wrong with the conversion.")
+    sys.exit(1)
+
+missing, unexpected = model.load_state_dict(filtered_state_dict, strict=False)
+if missing:
+    print(f"WARNING: {len(missing)} keys missing in converted weights (using base model weights)")
+    for k in missing[:10]:
+        print(f"  Missing: {k}")
+    if len(missing) > 10:
+        print(f"  ... and {len(missing) - 10} more")
+
+# 保存为 HF 格式（使用 safetensors，更安全、更快）
+print(f"Saving model to {output_dir}...")
+model.save_pretrained(output_dir, safe_serialization=True)
+print("Done!")
+
+# 清理临时文件
+print("Cleaning up temp files...")
+import shutil
+shutil.rmtree(temp_weights_dir)
+PYEOF
+
+export OUTPUT_DIR="$OUTPUT_DIR"
+export CHECKPOINT_DIR="$CHECKPOINT_DIR"
+export BASE_MODEL_DIR="$PROJECT_ROOT/a3b_ckpt"
+export PROJECT_ROOT="$PROJECT_ROOT"
+
+python3 /tmp/convert_to_hf.py
+
+echo ""
+echo "=========================================="
+echo "Conversion completed!"
+echo "HF format model saved to: $OUTPUT_DIR"
+echo "=========================================="

train/llama_factory_support/dataset_info.json

@@ -0,0 +1,44 @@
+{
+  "hy_v3_demo": {
+    "file_name": "../example_data.jsonl",
+    "formatting": "sharegpt",
+    "columns": {
+      "messages": "messages"
+    },
+    "tags": {
+      "role_tag": "role",
+      "content_tag": "content",
+      "user_tag": "user",
+      "assistant_tag": "assistant",
+      "system_tag": "system"
+    }
+  },
+  "hy_v3_translation": {
+    "file_name": "../data/example_data.jsonl",
+    "formatting": "sharegpt",
+    "columns": {
+      "messages": "messages"
+    },
+    "tags": {
+      "role_tag": "role",
+      "content_tag": "content",
+      "user_tag": "user",
+      "assistant_tag": "assistant",
+      "system_tag": "system"
+    }
+  },
+  "hy_dense_demo": {
+    "file_name": "../data/example_data.jsonl",
+    "formatting": "sharegpt",
+    "columns": {
+      "messages": "messages"
+    },
+    "tags": {
+      "role_tag": "role",
+      "content_tag": "content",
+      "user_tag": "user",
+      "assistant_tag": "assistant",
+      "system_tag": "system"
+    }
+  }
+}

train/llama_factory_support/ds_zero2_offload.json

@@ -0,0 +1,37 @@
+{
+    "fp16": {
+        "enabled": false,
+        "loss_scale": 0,
+        "loss_scale_window": 1000,
+        "initial_scale_power": 16,
+        "hysteresis": 2,
+        "min_loss_scale": 1
+    },
+    "bf16": {
+        "enabled": true
+    },
+
+    "zero_optimization": {
+        "stage": 2,
+        "offload_optimizer": {
+            "device": "cpu",
+            "pin_memory": true
+        },
+        "offload_param": {
+            "device": "cpu",
+            "pin_memory": true
+        },
+        "allgather_partitions": true,
+        "allgather_bucket_size": 5e8,
+        "overlap_comm": true,
+        "reduce_scatter": true,
+        "reduce_bucket_size": 5e8,
+        "contiguous_gradients": true
+    },
+
+    "gradient_accumulation_steps": "auto",
+    "gradient_clipping": "auto",
+    "steps_per_print": 10,
+    "train_micro_batch_size_per_gpu": "auto",
+    "wall_clock_breakdown": false
+}

train/llama_factory_support/ds_zero2_offload_lora.json

@@ -0,0 +1,33 @@
+{
+    "fp16": {
+        "enabled": false,
+        "loss_scale": 0,
+        "loss_scale_window": 1000,
+        "initial_scale_power": 16,
+        "hysteresis": 2,
+        "min_loss_scale": 1
+    },
+    "bf16": {
+        "enabled": true
+    },
+
+    "zero_optimization": {
+        "stage": 2,
+        "offload_optimizer": {
+            "device": "cpu",
+            "pin_memory": true
+        },
+        "allgather_partitions": true,
+        "allgather_bucket_size": 5e8,
+        "overlap_comm": true,
+        "reduce_scatter": true,
+        "reduce_bucket_size": 5e8,
+        "contiguous_gradients": true
+    },
+
+    "gradient_accumulation_steps": "auto",
+    "gradient_clipping": "auto",
+    "steps_per_print": 10,
+    "train_micro_batch_size_per_gpu": "auto",
+    "wall_clock_breakdown": false
+}

train/llama_factory_support/ds_zero3_offload.json

@@ -0,0 +1,36 @@
+{
+    "fp16": {
+        "enabled": false,
+        "loss_scale": 0,
+        "loss_scale_window": 1000,
+        "initial_scale_power": 16,
+        "hysteresis": 2,
+        "min_loss_scale": 1
+    },
+    "bf16": {
+        "enabled": true
+    },
+
+    "zero_optimization": {
+        "stage": 3,
+        "offload_optimizer": {
+            "device": "cpu",
+            "pin_memory": true
+        },
+        "overlap_comm": true,
+        "contiguous_gradients": true,
+        "sub_group_size": 1e9,
+        "reduce_bucket_size": 1e8,
+        "stage3_prefetch_bucket_size": 1e8,
+        "stage3_param_persistence_threshold": 0,
+        "stage3_max_live_parameters": 1e8,
+        "stage3_max_reuse_distance": 1e8,
+        "stage3_gather_16bit_weights_on_model_save": true
+    },
+
+    "gradient_accumulation_steps": "auto",
+    "gradient_clipping": "auto",
+    "steps_per_print": 10,
+    "train_micro_batch_size_per_gpu": "auto",
+    "wall_clock_breakdown": false
+}

train/llama_factory_support/hy_dense_1_8b_full_sft.yaml

@@ -0,0 +1,50 @@
+### model
+model_name_or_path: path_to_dense_1_8b_model
+trust_remote_code: true
+
+### method
+stage: sft
+do_train: true
+finetuning_type: full
+deepspeed: ds_zero2_offload.json
+
+### dataset
+dataset_dir: .
+dataset: hy_dense_demo
+template: hy_dense_1_8b
+cutoff_len: 4096
+max_samples: 1000
+overwrite_cache: true
+
+### output
+output_dir: saves/hy_dense_1_8b/full/sft
+logging_steps: 1
+save_steps: 30
+plot_loss: true
+overwrite_output_dir: true
+save_only_model: false
+report_to: none  # choices: [none, wandb, tensorboard, swanlab, mlflow]
+
+### train
+per_device_train_batch_size: 1
+gradient_accumulation_steps: 1
+learning_rate: 1.0e-5
+num_train_epochs: 1.0
+lr_scheduler_type: cosine_with_min_lr
+lr_scheduler_kwargs:
+  min_lr_rate: 0.1  # min_lr / learning_rate = 1e-6 / 1e-5 = 0.1
+warmup_ratio: 0.1
+bf16: true
+gradient_checkpointing: true
+gradient_checkpointing_kwargs:
+  use_reentrant: true
+ddp_timeout: 180000000
+flash_attn: fa2
+resume_from_checkpoint: null
+
+### eval
+# eval_dataset: alpaca_en_demo
+# val_size: 0.1
+# per_device_eval_batch_size: 1
+# eval_strategy: steps
+# eval_steps: 500

train/llama_factory_support/hy_dense_1_8b_lora_sft.yaml

@@ -0,0 +1,56 @@
+### model
+model_name_or_path: path_to_dense_1_8b_model
+trust_remote_code: true
+
+### method
+stage: sft
+do_train: true
+finetuning_type: lora
+deepspeed: ds_zero2_offload_lora.json
+
+### LoRA parameters
+lora_rank: 64
+lora_alpha: 128
+lora_dropout: 0.05
+lora_target: q_proj,k_proj,v_proj,o_proj
+
+### dataset
+dataset_dir: .
+dataset: hy_dense_demo
+template: hy_dense_1_8b
+cutoff_len: 4096
+max_samples: 1000
+overwrite_cache: true
+
+### output
+output_dir: saves/hy_dense_1_8b/lora/sft
+logging_steps: 1
+save_steps: 10
+plot_loss: true
+overwrite_output_dir: true
+save_only_model: false
+report_to: none  # choices: [none, wandb, tensorboard, swanlab, mlflow]
+
+### train
+per_device_train_batch_size: 1
+gradient_accumulation_steps: 1
+learning_rate: 2.0e-4
+num_train_epochs: 1.0
+lr_scheduler_type: cosine_with_min_lr
+lr_scheduler_kwargs:
+  min_lr_rate: 0.1  # min_lr / learning_rate = 2e-5 / 2e-4 = 0.1
+warmup_ratio: 0.1
+bf16: true
+gradient_checkpointing: true
+gradient_checkpointing_kwargs:
+  use_reentrant: true
+ddp_timeout: 180000000
+flash_attn: fa2
+resume_from_checkpoint: null
+
+### eval
+# eval_dataset: alpaca_en_demo
+# val_size: 0.1
+# per_device_eval_batch_size: 1
+# eval_strategy: steps
+# eval_steps: 500

train/llama_factory_support/hy_dense_7b_full_sft.yaml

@@ -0,0 +1,50 @@
+### model
+model_name_or_path: path_to_dense_7b_model
+trust_remote_code: true
+
+### method
+stage: sft
+do_train: true
+finetuning_type: full
+deepspeed: ds_zero3_offload.json
+
+### dataset
+dataset_dir: .
+dataset: hy_dense_demo
+template: hy_dense_7b
+cutoff_len: 4096
+max_samples: 1000
+overwrite_cache: true
+
+### output
+output_dir: saves/hy_dense_7b/full/sft
+logging_steps: 1
+save_steps: 10
+plot_loss: true
+overwrite_output_dir: true
+save_only_model: false
+report_to: none  # choices: [none, wandb, tensorboard, swanlab, mlflow]
+
+### train
+per_device_train_batch_size: 1
+gradient_accumulation_steps: 1
+learning_rate: 1.0e-5
+num_train_epochs: 1.0
+lr_scheduler_type: cosine_with_min_lr
+lr_scheduler_kwargs:
+  min_lr_rate: 0.1  # min_lr / learning_rate = 1e-6 / 1e-5 = 0.1
+warmup_ratio: 0.1
+bf16: true
+gradient_checkpointing: true
+gradient_checkpointing_kwargs:
+  use_reentrant: true
+ddp_timeout: 180000000
+flash_attn: fa2
+resume_from_checkpoint: null
+
+### eval
+# eval_dataset: alpaca_en_demo
+# val_size: 0.1
+# per_device_eval_batch_size: 1
+# eval_strategy: steps
+# eval_steps: 500

train/llama_factory_support/hy_dense_7b_lora_sft.yaml

@@ -0,0 +1,56 @@
+### model
+model_name_or_path: path_to_dense_7b_model
+trust_remote_code: true
+
+### method
+stage: sft
+do_train: true
+finetuning_type: lora
+deepspeed: ds_zero2_offload_lora.json
+
+### LoRA parameters
+lora_rank: 64
+lora_alpha: 128
+lora_dropout: 0.05
+lora_target: q_proj,k_proj,v_proj,o_proj
+
+### dataset
+dataset_dir: .
+dataset: hy_dense_demo
+template: hy_dense_7b
+cutoff_len: 4096
+max_samples: 1000
+overwrite_cache: true
+
+### output
+output_dir: saves/hy_dense_7b/lora/sft
+logging_steps: 1
+save_steps: 30
+plot_loss: true
+overwrite_output_dir: true
+save_only_model: false
+report_to: none  # choices: [none, wandb, tensorboard, swanlab, mlflow]
+
+### train
+per_device_train_batch_size: 1
+gradient_accumulation_steps: 1
+learning_rate: 2.0e-4
+num_train_epochs: 1.0
+lr_scheduler_type: cosine_with_min_lr
+lr_scheduler_kwargs:
+  min_lr_rate: 0.1  # min_lr / learning_rate = 2e-5 / 2e-4 = 0.1
+warmup_ratio: 0.1
+bf16: true
+gradient_checkpointing: true
+gradient_checkpointing_kwargs:
+  use_reentrant: true
+ddp_timeout: 180000000
+flash_attn: fa2
+resume_from_checkpoint: null
+
+### eval
+# eval_dataset: alpaca_en_demo
+# val_size: 0.1
+# per_device_eval_batch_size: 1
+# eval_strategy: steps
+# eval_steps: 500

train/llama_factory_support/hy_dense_template.py

@@ -0,0 +1,82 @@
+"""
+HunYuan Dense model chat template registration for LLaMA Factory.
+
+Registers two templates:
+  - hy_dense_1_8b: for HunYuan Dense 1.8B model (and 0.5B/4B)
+  - hy_dense_7b:   for HunYuan Dense 7B model
+
+Usage:
+    1. Copy this file's register_template blocks into LLaMA Factory's
+       src/llamafactory/data/template.py  (for upstream MR).
+    2. Or import this module before training to register at runtime:
+       import hy_dense_template
+
+Note:
+    The existing LLaMA Factory built-in templates `hunyuan` and `hunyuan_small`
+    have subtle differences from the official chat_template.jinja files shipped
+    with the models. These new templates are designed to match the official
+    jinja templates exactly.
+"""
+
+from llamafactory.data.template import register_template
+from llamafactory.data.formatter import EmptyFormatter, StringFormatter
+
+# ---------------------------------------------------------------------------
+# Dense 1.8B chat template (also applies to 0.5B/4B)
+#
+# Token format (from dense_1_8b_0508/global_step_560/chat_template.jinja):
+#   BOS:        <｜hy_begin▁of▁sentence｜>
+#   System:     {system_content}<｜hy_place▁holder▁no▁3｜>
+#   User:       <｜hy_User｜>{user_content}
+#   Assistant:  <｜hy_Assistant｜>{assistant_content}<｜hy_place▁holder▁no▁2｜>
+#   Stop:       <｜hy_place▁holder▁no▁2｜>
+#
+# Key differences from LF built-in `hunyuan_small`:
+#   - User format: NO trailing <｜hy_place▁holder▁no▁8｜> after user content
+#   - Assistant format: HAS <｜hy_Assistant｜> prefix before assistant content
+#
+# The eos_token in tokenizer_config.json is <｜hy_place▁holder▁no▁2｜>,
+# so we use efficient_eos=True to let LF append it via {eos_token} slot.
+# ---------------------------------------------------------------------------
+
+register_template(
+    name="hy_dense_1_8b",
+    format_user=StringFormatter(slots=["<｜hy_User｜>{{content}}"]),
+    format_assistant=StringFormatter(slots=["<｜hy_Assistant｜>{{content}}", {"eos_token"}]),
+    format_system=StringFormatter(slots=["{{content}}<｜hy_place▁holder▁no▁3｜>"]),
+    format_prefix=EmptyFormatter(slots=[{"bos_token"}]),
+    stop_words=["<｜hy_place▁holder▁no▁2｜>"],
+    efficient_eos=True,
+)
+
+
+# ---------------------------------------------------------------------------
+# Dense 7B chat template
+#
+# Token format (from dense_7b_0509/global_step_560/chat_template.jinja):
+#   BOS:        <|startoftext|>
+#   System:     {system_content}<|extra_4|>
+#   User:       {user_content}<|extra_0|>
+#   Assistant:  {assistant_content}<|eos|>
+#   Stop:       <|eos|>
+#
+# Key differences from LF built-in `hunyuan`:
+#   - Uses {bos_token} and {eos_token} slots for portability
+#   - efficient_eos=True to use tokenizer's eos_token
+#
+# Note on multi-turn: The official jinja adds <|startoftext|> before each
+# user message (except the first one when system is present). LLaMA Factory's
+# format_prefix only adds BOS once at the beginning. For single-turn training
+# this is correct. For multi-turn, there is a minor discrepancy (missing
+# <|startoftext|> before 2nd+ user turns), which is acceptable for fine-tuning.
+# ---------------------------------------------------------------------------
+
+register_template(
+    name="hy_dense_7b",
+    format_user=StringFormatter(slots=["{{content}}<|extra_0|>"]),
+    format_assistant=StringFormatter(slots=["{{content}}", {"eos_token"}]),
+    format_system=StringFormatter(slots=["{{content}}<|extra_4|>"]),
+    format_prefix=EmptyFormatter(slots=[{"bos_token"}]),
+    stop_words=["<|eos|>"],
+    efficient_eos=True,
+)

train/llama_factory_support/hy_v3_full_sft.yaml

@@ -0,0 +1,50 @@
+### model
+model_name_or_path: ../hf
+trust_remote_code: true
+
+### method
+stage: sft
+do_train: true
+finetuning_type: full
+deepspeed: ds_zero3_offload.json
+
+### dataset
+dataset_dir: .
+dataset: hy_v3_demo
+template: hy_v3
+cutoff_len: 4096  # HYV3 supports 262k context length
+max_samples: 1000
+overwrite_cache: true
+
+### output
+output_dir: saves/hy_v3/full/sft
+logging_steps: 1
+save_steps: 10
+plot_loss: true
+overwrite_output_dir: true
+save_only_model: false
+report_to: none  # choices: [none, wandb, tensorboard, swanlab, mlflow]
+
+### train
+per_device_train_batch_size: 1
+gradient_accumulation_steps: 1
+learning_rate: 1.0e-5
+num_train_epochs: 3.0
+lr_scheduler_type: cosine_with_min_lr
+lr_scheduler_kwargs:
+  min_lr_rate: 0.1  # min_lr / learning_rate = 1e-6 / 1e-5 = 0.1
+warmup_ratio: 0.1
+bf16: true
+gradient_checkpointing: true
+gradient_checkpointing_kwargs:
+  use_reentrant: true
+ddp_timeout: 180000000
+flash_attn: fa2
+resume_from_checkpoint: null
+
+### eval
+# eval_dataset: alpaca_en_demo
+# val_size: 0.1
+# per_device_eval_batch_size: 1
+# eval_strategy: steps
+# eval_steps: 500

train/llama_factory_support/hy_v3_lora_sft.yaml

@@ -0,0 +1,56 @@
+### model
+model_name_or_path: ../hf
+trust_remote_code: true
+
+### method
+stage: sft
+do_train: true
+finetuning_type: lora
+deepspeed: ds_zero2_offload_lora.json
+
+### LoRA parameters
+lora_rank: 64
+lora_alpha: 128
+lora_dropout: 0.05
+lora_target: q_proj,k_proj,v_proj,o_proj
+
+### dataset
+dataset_dir: .
+dataset: hy_v3_translation
+template: hy_v3
+cutoff_len: 4096  # Use shorter context for LoRA to save memory; increase if needed
+max_samples: 1000
+overwrite_cache: true
+
+### output
+output_dir: saves/hy_v3/lora/sft
+logging_steps: 10
+save_steps: 500
+plot_loss: true
+overwrite_output_dir: true
+save_only_model: false
+report_to: none  # choices: [none, wandb, tensorboard, swanlab, mlflow]
+
+### train
+per_device_train_batch_size: 1
+gradient_accumulation_steps: 1
+learning_rate: 2.0e-4
+num_train_epochs: 3.0
+lr_scheduler_type: cosine_with_min_lr
+lr_scheduler_kwargs:
+  min_lr_rate: 0.1  # min_lr / learning_rate = 2e-5 / 2e-4 = 0.1
+warmup_ratio: 0.1
+bf16: true
+gradient_checkpointing: true
+gradient_checkpointing_kwargs:
+  use_reentrant: true
+ddp_timeout: 180000000
+flash_attn: fa2
+resume_from_checkpoint: null
+
+### eval
+# eval_dataset: alpaca_en_demo
+# val_size: 0.1
+# per_device_eval_batch_size: 1
+# eval_strategy: steps
+# eval_steps: 500

train/llama_factory_support/hy_v3_patches.py

@@ -0,0 +1,215 @@
+"""
+HYV3 monkey-patches for LLaMA Factory + DeepSpeed training.
+
+This module applies all necessary runtime patches so that HYV3 (MoE)
+can be trained correctly under LLaMA Factory with DeepSpeed.
+
+Usage:
+    Import this module **before** calling `llamafactory-cli train`:
+
+        import hy_v3_patches          # applies patches on import
+        # ... then start training
+
+    Or add to the LLaMA Factory YAML via a custom entry-point wrapper.
+
+Patches applied:
+    1. (Removed) -- transformers 5.8.1+ has built-in conversion_mapping for
+       hy_v3 that handles key renaming + expert fusing automatically.
+    2. Router forward dtype fix (MoE router gate dtype alignment for ZeRO-3)
+    3. gradient_checkpointing   (use_reentrant=True for ZeRO-3)
+    4. Tokenizer file copy      (CustomSaveCallback)
+    5. (Removed) -- was per-expert ModuleList, now using native 3D Parameters
+    6. (Removed) -- transformers 5.8.1+ has built-in revert_weight_conversion
+       in save_pretrained that handles outer->inner format automatically.
+"""
+
+import os
+import logging
+import shutil
+from typing import Optional
+
+import torch
+import torch.nn.functional as _F
+
+logger = logging.getLogger(__name__)
+
+# ============================================================================
+# Patch 2: Router forward dtype alignment for ZeRO-3
+#
+# The HYV3 MoE HYV3TopKRouter.forward() calls F.linear with .float().
+# Under DeepSpeed ZeRO-3, F.linear is replaced by zero3_linear_wrap which
+# internally does input.matmul(weight.t()) WITHOUT aligning dtypes.
+# When ZeRO-3 stores the gate weight in bf16, the fp32 input causes a
+# dtype mismatch RuntimeError.
+#
+# Fix: monkey-patch HYV3TopKRouter.forward to cast input to
+# self.weight.dtype before F.linear, then cast the output back to float32.
+# ============================================================================
+
+_router_patch_applied = False
+
+def _apply_router_dtype_patch():
+    """Monkey-patch HYV3TopKRouter.forward to align gate input dtype with weight dtype."""
+    global _router_patch_applied
+    if _router_patch_applied:
+        return
+
+    try:
+        from transformers.models.hy_v3.modeling_hy_v3 import HYV3TopKRouter
+    except ImportError:
+        try:
+            from transformers.hy_v3.modeling_hy_v3 import HYV3TopKRouter
+        except ImportError:
+            logger.warning(
+                "Could not import HYV3TopKRouter; "
+                "router dtype patch NOT applied."
+            )
+            return
+
+    def _patched_router_forward(
+        self,
+        hidden_states: torch.Tensor,
+        e_score_correction_bias: torch.Tensor,
+    ) -> tuple:
+        hidden_states = hidden_states.reshape(-1, self.hidden_dim)
+        # Cast input to match weight dtype (bf16 under ZeRO-3)
+        # instead of hard-coding float32, to avoid matmul dtype mismatch.
+        weight_dtype = self.weight.dtype
+        router_logits = _F.linear(hidden_states.to(weight_dtype), self.weight.to(weight_dtype))
+        # Cast back to float32 for numerically stable sigmoid
+        router_logits = router_logits.to(torch.float32)
+        routing_weights = torch.sigmoid(router_logits)
+
+        scores_for_choice = routing_weights + e_score_correction_bias
+        _, top_k_index = torch.topk(scores_for_choice, self.top_k, dim=-1, sorted=False)
+        top_k_weights = routing_weights.gather(1, top_k_index)
+
+        top_k_weights = top_k_weights / (top_k_weights.sum(dim=-1, keepdim=True) + 1e-20)
+        top_k_weights = top_k_weights * self.router_scaling_factor
+
+        return router_logits, top_k_weights, top_k_index
+
+    HYV3TopKRouter.forward = _patched_router_forward
+    _router_patch_applied = True
+    logger.info("HYV3 patch applied: HYV3TopKRouter.forward dtype alignment for ZeRO-3.")
+
+# ============================================================================
+# Patch 3: gradient_checkpointing use_reentrant=True
+#
+# PyTorch's torch.utils.checkpoint with use_reentrant=False (the default
+# in transformers) performs strict metadata checks on recomputed tensors.
+# Under ZeRO-3, parameters are all-gathered during the first forward pass
+# but may be partitioned back when the checkpoint recomputes, causing a
+# CheckpointError.  Setting use_reentrant=True avoids this.
+#
+# This is applied via a Trainer callback that modifies training_args
+# before training starts.
+# ============================================================================
+
+# ============================================================================
+# Patch 4: Tokenizer file copy callback
+#
+# Ensures each checkpoint directory is self-contained for inference by
+# copying all tokenizer-related files from the original tokenizer path.
+# ============================================================================
+
+# Tokenizer files that should be copied to each checkpoint
+_TOKENIZER_FILES = [
+    "generation_config.json",
+    "hy.tiktoken",
+    "tokenizer_config.json",
+    "tokenization_hy.py",
+    "tokenizer.json",
+    "special_tokens_map.json",
+    "chat_template.jinja",
+]
+
+def _copy_tokenizer_to_checkpoint(tokenizer_dir: str, checkpoint_dir: str):
+    """Copy tokenizer files from tokenizer_dir to checkpoint_dir."""
+    for fname in _TOKENIZER_FILES:
+        src = os.path.join(tokenizer_dir, fname)
+        if os.path.isfile(src):
+            shutil.copy(src, os.path.join(checkpoint_dir, fname))
+
+# ============================================================================
+# LLaMA Factory Callback: integrates patches 3, 4 into the training loop
+# ============================================================================
+
+try:
+    from transformers import TrainerCallback
+    from transformers.trainer_utils import PREFIX_CHECKPOINT_DIR
+
+    class HYV3PatchCallback(TrainerCallback):
+        """
+        LLaMA Factory compatible callback that applies HYV3-specific patches.
+
+        Add to your YAML or pass to Trainer:
+            callbacks: [hy_v3_patches.HYV3PatchCallback]
+        """
+
+        def __init__(self, tokenizer_dir: Optional[str] = None):
+            """
+            Args:
+                tokenizer_dir: Path to the original tokenizer directory.
+                    If None, will try to use model_name_or_path from training args.
+            """
+            self._tokenizer_dir = tokenizer_dir
+
+        def on_train_begin(self, args, state, control, **kwargs):
+            # --- Patch 3: gradient_checkpointing use_reentrant ---
+            if getattr(args, "gradient_checkpointing", False) and getattr(args, "deepspeed", None):
+                if not hasattr(args, "gradient_checkpointing_kwargs") or not args.gradient_checkpointing_kwargs:
+                    args.gradient_checkpointing_kwargs = {"use_reentrant": True}
+                elif "use_reentrant" not in args.gradient_checkpointing_kwargs:
+                    args.gradient_checkpointing_kwargs["use_reentrant"] = True
+                logger.info("HYV3 patch applied: gradient_checkpointing use_reentrant=True.")
+
+            return control
+
+        def on_save(self, args, state, control, **kwargs):
+            # --- Patch 4: Copy tokenizer files ---
+            if torch.distributed.is_initialized() and torch.distributed.get_rank() != 0:
+                return control
+
+            checkpoint_dir = os.path.join(
+                args.output_dir,
+                f"{PREFIX_CHECKPOINT_DIR}-{state.global_step}",
+            )
+
+            # Determine tokenizer directory
+            tokenizer_dir = self._tokenizer_dir
+            if tokenizer_dir is None:
+                # Try common locations
+                tokenizer_dir = getattr(args, "tokenizer_name_or_path", None)
+                if tokenizer_dir is None:
+                    tokenizer_dir = getattr(args, "model_name_or_path", None)
+
+            if tokenizer_dir and os.path.isdir(tokenizer_dir):
+                _copy_tokenizer_to_checkpoint(tokenizer_dir, checkpoint_dir)
+                logger.info(
+                    "HYV3: Copied tokenizer files from %s to %s",
+                    tokenizer_dir, checkpoint_dir
+                )
+
+            return control
+
+except ImportError:
+    logger.warning(
+        "transformers not available; HYV3PatchCallback not defined."
+    )
+
+# ============================================================================
+# Auto-apply patches on import
+# ============================================================================
+
+# Patch 2: Router dtype fix
+_apply_router_dtype_patch()
+
+# Patches 3, 4 are applied via HYV3PatchCallback during training.
+# Users should add HYV3PatchCallback to their Trainer callbacks.
+
+logger.info(
+    "HYV3 patches module loaded. Patch 2 (Router dtype fix) applied. "
+    "Remember to add HYV3PatchCallback to your Trainer callbacks "
+    "for gradient_checkpointing and tokenizer copy support."
+)

train/llama_factory_support/hy_v3_template.py

@@ -0,0 +1,46 @@
+"""
+HYV3 chat template registration for LLaMA Factory.
+
+Usage:
+    1. Copy this file's register_template block into LLaMA Factory's
+       src/llamafactory/data/template.py  (for upstream MR).
+    2. Or import this module before training to register at runtime:
+       import hy_v3_template
+"""
+
+from llamafactory.data.template import ReasoningTemplate, register_template
+from llamafactory.data.formatter import EmptyFormatter, StringFormatter
+
+
+# ---------------------------------------------------------------------------
+# HYV3 (MoE, pure text) chat template
+#
+# Token format (from chat_template.jinja & tokenizer_config.json):
+#   BOS:        <｜hy_begin▁of▁sentence｜>
+#   System:    {system_content}                (directly after BOS, no role tag)
+#   User:       <｜hy_User｜>{user_content}
+#   Assistant:  <｜hy_Assistant｜>{assistant_content}<｜hy_eos｜>
+#   EOS:       <｜hy_eos｜>
+#
+# Loss mask: only compute loss on assistant content (including <｜hy_eos｜>).
+#
+# Note: The system message has NO explicit role token -- it is placed right
+# after BOS.  The eos_token is <｜hy_eos｜>.
+#
+# Reasoning: Supports think tags via ReasoningTemplate.
+#   - thought_words: ("<think>", "</think>") matching jinja template
+#   - enable_thinking: set globally via data_args.enable_thinking (default True)
+#   - Training data always includes think tags (empty or with content)
+# ---------------------------------------------------------------------------
+
+register_template(
+    name="hy_v3",
+    template_class=ReasoningTemplate,
+    format_user=StringFormatter(slots=["<｜hy_User｜>{{content}}"]),
+    format_assistant=StringFormatter(slots=["<｜hy_Assistant｜>{{content}}", {"eos_token"}]),
+    format_system=StringFormatter(slots=["{{content}}"]),
+    format_prefix=EmptyFormatter(slots=[{"bos_token"}]),
+    thought_words=("<think>", "</think>"),
+    stop_words=["<｜hy_eos｜>"],
+    efficient_eos=True,
+)

train/llama_factory_support/train_hy_dense.py

@@ -0,0 +1,81 @@
+"""
+LLaMA Factory training entry-point wrapper for HunYuan Dense models.
+
+This script:
+  1. Registers the hy_dense_1_8b and hy_dense_7b chat templates
+  2. Injects a lightweight PatchCallback (tokenizer copy + gradient checkpointing fix)
+  3. Calls run_exp() to start LLaMA Factory training
+
+How it works:
+  - train_lf_dense.sh launches this script via torchrun directly:
+        torchrun ... train_hy_dense.py hy_dense_1_8b_full_sft.yaml
+  - Each torchrun worker executes this script, so all patches are applied
+    in every worker process before training begins.
+  - We call run_exp() directly (not the CLI launcher) to avoid the
+    launcher re-spawning workers and losing our patches.
+
+Note:
+  Dense models do NOT need MoE-specific patches (router dtype fix, expert
+  key rename, etc.). Only the tokenizer copy callback and gradient
+  checkpointing fix are needed.
+
+Usage:
+    # Via launch script (recommended):
+    bash train_lf_dense.sh
+
+    # Direct single-node (1 GPU, 1.8B model):
+    torchrun --nproc_per_node 1 train_hy_dense.py hy_dense_1_8b_full_sft.yaml
+
+    # Direct single-node (2 GPUs, 7B model):
+    torchrun --nproc_per_node 2 train_hy_dense.py hy_dense_7b_full_sft.yaml
+"""
+
+import sys
+import os
+
+# Add current directory to path so templates can be imported
+sys.path.insert(0, os.path.dirname(os.path.abspath(__file__)))
+
+# Step 1: Register Dense model templates (must be before training starts)
+import hy_dense_template  # noqa: F401
+
+# Step 2: Import the patch callback (reuse HYV3PatchCallback for tokenizer copy)
+# The MoE router patch will be silently skipped since Dense models don't have
+# HYV3TopKRouter. Only Patch 3 (gradient_checkpointing) and Patch 4 (tokenizer
+# copy) will be effective.
+import hy_v3_patches  # noqa: F401
+
+# Step 3: Inject PatchCallback into LLaMA Factory's training flow
+from llamafactory.train.sft.workflow import run_sft as _orig_run_sft
+
+
+def _patched_run_sft(model_args, data_args, training_args, finetuning_args, generating_args, callbacks=None):
+    """Wrap run_sft to inject HYV3PatchCallback for tokenizer copy."""
+    if callbacks is None:
+        callbacks = []
+
+    # Determine tokenizer directory for the save callback
+    tokenizer_dir = getattr(model_args, "model_name_or_path", None)
+    callbacks.append(hy_v3_patches.HYV3PatchCallback(tokenizer_dir=tokenizer_dir))
+
+    return _orig_run_sft(model_args, data_args, training_args, finetuning_args, generating_args, callbacks=callbacks)
+
+
+# Monkey-patch the SFT workflow
+import llamafactory.train.sft.workflow as _sft_wf
+_sft_wf.run_sft = _patched_run_sft
+
+
+def main():
+    """Entry point: called by torchrun in each worker process.
+
+    Since train_lf_dense.sh launches us via torchrun directly, all patches
+    (template registration, tokenizer copy callback injection) are already
+    applied in this process. We just call run_exp() to start training.
+    """
+    from llamafactory.train.tuner import run_exp
+    run_exp()
+
+
+if __name__ == "__main__":
+    main()

train/llama_factory_support/train_hy_v3.py

@@ -0,0 +1,72 @@
+"""
+LLaMA Factory training entry-point wrapper for HYV3.
+
+This script:
+  1. Registers the hy_v3 chat template
+  2. Applies all HYV3 monkey-patches (checkpoint key rename, dtype fix, etc.)
+  3. Injects HYV3PatchCallback into the training loop
+  4. Calls run_exp() to start LLaMA Factory training
+
+How it works:
+  - train_lf.sh launches this script via torchrun directly:
+        torchrun ... train_hy_v3.py hy_v3_full_sft.yaml
+  - Each torchrun worker executes this script, so all patches are applied
+    in every worker process before training begins.
+  - We call run_exp() directly (not the CLI launcher) to avoid the
+    launcher re-spawning workers and losing our patches.
+
+Usage:
+    # Via launch script (recommended):
+    bash train_lf.sh
+
+    # Direct single-node (8 GPUs):
+    torchrun --nproc_per_node 8 train_hy_v3.py hy_v3_full_sft.yaml
+"""
+
+import sys
+import os
+
+# Add current directory to path so patches can be imported
+sys.path.insert(0, os.path.dirname(os.path.abspath(__file__)))
+
+# Step 1: Register HYV3 template (must be before training starts)
+import hy_v3_template  # noqa: F401
+
+# Step 2: Apply checkpoint key rename patch (must be before model loading)
+import hy_v3_patches  # noqa: F401
+
+# Step 3: Inject HYV3PatchCallback into LLaMA Factory's training flow
+from llamafactory.train.sft.workflow import run_sft as _orig_run_sft
+
+
+def _patched_run_sft(model_args, data_args, training_args, finetuning_args, generating_args, callbacks=None):
+    """Wrap run_sft to inject HYV3PatchCallback."""
+    if callbacks is None:
+        callbacks = []
+
+    # Determine tokenizer directory for the save callback
+    tokenizer_dir = getattr(model_args, "model_name_or_path", None)
+    callbacks.append(hy_v3_patches.HYV3PatchCallback(tokenizer_dir=tokenizer_dir))
+
+    return _orig_run_sft(model_args, data_args, training_args, finetuning_args, generating_args, callbacks=callbacks)
+
+
+# Monkey-patch the SFT workflow
+import llamafactory.train.sft.workflow as _sft_wf
+_sft_wf.run_sft = _patched_run_sft
+
+
+def main():
+    """Entry point: called by torchrun in each worker process.
+
+    Since train_lf.sh launches us via torchrun directly, all patches
+    (template registration, checkpoint key rename, SFT callback injection)
+    are already applied in this process.  We just call run_exp() to start
+    training — no need to go through the CLI launcher.
+    """
+    from llamafactory.train.tuner import run_exp
+    run_exp()
+
+
+if __name__ == "__main__":
+    main()

train/llama_factory_support/train_lf.sh

@@ -0,0 +1,114 @@
+#!/bin/bash
+# ============================================================================
+# LLaMA Factory training launch script for HYV3
+#
+# This script sets up the environment and launches training via torchrun.
+#
+# We use train_hy_v3.py as the entry point (not llamafactory-cli)
+# because we need to inject HYV3-specific monkey-patches and register
+# the hy_v3 chat template BEFORE LLaMA Factory starts.
+# train_hy_v3.py directly calls run_exp() in each torchrun worker,
+# ensuring all patches are active.
+#
+# Usage:
+#   Single node:  bash train_lf.sh
+#   Multi-node:   Run this script on EACH node with the same IP_LIST.
+#                 IP_LIST="10.0.0.1,10.0.0.2" bash train_lf.sh
+# ============================================================================
+
+set -euo pipefail
+
+# -------------------- Network Configuration --------------------
+NET_TYPE="high"
+export NCCL_DEBUG=WARN
+export NCCL_P2P_LEVEL=NVL
+export NCCL_IB_TIMEOUT=24
+export NCCL_NVLS_ENABLE=0
+export NCCL_MPI_PROFILE_PRIMS_ENABLE=0
+export CUDA_DEVICE_MAX_CONNECTIONS=1
+export TORCH_NCCL_HEARTBEAT_TIMEOUT_SEC=3600
+if [[ "${NET_TYPE}" = "low" ]]; then
+    export NCCL_SOCKET_IFNAME=eth1
+    export NCCL_IB_GID_INDEX=3
+    export NCCL_IB_HCA=mlx5_2:1
+    export NCCL_IB_SL=3
+    export NCCL_CHECK_DISABLE=1
+    export NCCL_P2P_DISABLE=0
+    export NCCL_LL_THRESHOLD=16384
+    export NCCL_IB_CUDA_SUPPORT=1
+else
+    export NCCL_IB_GID_INDEX=3
+    export NCCL_IB_SL=3
+    export NCCL_CHECK_DISABLE=1
+    export NCCL_P2P_DISABLE=0
+    export NCCL_IB_DISABLE=0
+    export NCCL_LL_THRESHOLD=16384
+    export NCCL_IB_CUDA_SUPPORT=1
+    export NCCL_SOCKET_IFNAME=bond1
+    export UCX_NET_DEVICES=bond1
+    export NCCL_IB_HCA=mlx5_bond_1,mlx5_bond_5,mlx5_bond_3,mlx5_bond_7,mlx5_bond_4,mlx5_bond_8,mlx5_bond_2,mlx5_bond_6
+    export NCCL_COLLNET_ENABLE=0
+    export SHARP_COLL_ENABLE_SAT=0
+    export NCCL_NET_GDR_LEVEL=2
+    export NCCL_IB_QPS_PER_CONNECTION=4
+    export NCCL_IB_TC=160
+    export NCCL_PXN_DISABLE=1
+fi
+
+# Skip LLaMA Factory version check (we use a newer transformers branch)
+export DISABLE_VERSION_CHECK=1
+
+# -------------------- Node Configuration --------------------
+export HOST_GPU_NUM=8
+# IP list, comma separated. e.g. "10.0.0.1,10.0.0.2" or single node "127.0.0.1"
+export IP_LIST=${IP_LIST:-"127.0.0.1"}
+
+MASTER_PORT=${MASTER_PORT:-29500}
+
+IFS=',' read -ra IP_ARRAY <<< "$IP_LIST"
+NODES=${#IP_ARRAY[@]}
+MASTER_ADDR=${IP_ARRAY[0]}
+
+# -------------------- Paths --------------------
+SCRIPT_DIR="$(cd "$(dirname "${BASH_SOURCE[0]}")" && pwd)"
+YAML_FILE="${YAML_FILE:-${SCRIPT_DIR}/hy_v3_full_sft.yaml}"
+ENTRY_SCRIPT="${SCRIPT_DIR}/train_hy_v3.py"
+
+# -------------------- Distributed Environment --------------------
+export MASTER_ADDR="${MASTER_ADDR}"
+export MASTER_PORT="${MASTER_PORT}"
+export NNODES="${NODES}"
+
+if [ ${NODES} -gt 1 ]; then
+    # Determine local node rank by matching local IP against IP_LIST
+    LOCAL_IP=$(hostname -i | awk '{print $1}')
+    NODE_RANK=0
+    for i in "${!IP_ARRAY[@]}"; do
+        if [[ "${IP_ARRAY[$i]}" == "${LOCAL_IP}" ]]; then
+            NODE_RANK=$i
+            break
+        fi
+    done
+    export RANK="${NODE_RANK}"
+else
+    export RANK=0
+fi
+
+echo "============================================"
+echo "  HYV3 LLaMA Factory Training"
+echo "  Nodes: ${NNODES}, Rank: ${RANK}"
+echo "  Master: ${MASTER_ADDR}:${MASTER_PORT}"
+echo "  GPUs per node: ${HOST_GPU_NUM}"
+echo "  Total GPUs: $((NODES * HOST_GPU_NUM))"
+echo "============================================"
+
+# -------------------- Launch --------------------
+# We launch torchrun directly (instead of FORCE_TORCHRUN) so that each
+# worker process runs train_hy_v3.py with all HYV3 patches applied.
+torchrun \
+    --nnodes "${NNODES}" \
+    --node_rank "${RANK}" \
+    --nproc_per_node "${HOST_GPU_NUM}" \
+    --master_addr "${MASTER_ADDR}" \
+    --master_port "${MASTER_PORT}" \
+    "${ENTRY_SCRIPT}" "${YAML_FILE}"

train/llama_factory_support/train_lf_dense.sh

@@ -0,0 +1,124 @@
+#!/bin/bash
+# ============================================================================
+# LLaMA Factory training launch script for HunYuan Dense models (1.8B / 7B)
+#
+# This script sets up the environment and launches training via torchrun.
+#
+# We use train_hy_dense.py as the entry point (not llamafactory-cli)
+# because we need to register the hy_dense chat templates and inject
+# the PatchCallback BEFORE LLaMA Factory starts.
+# train_hy_dense.py directly calls run_exp() in each torchrun worker,
+# ensuring all patches are active.
+#
+# Usage:
+#   Single node (1.8B, default):
+#     bash train_lf_dense.sh
+#
+#   Single node (7B):
+#     YAML_FILE=hy_dense_7b_full_sft.yaml bash train_lf_dense.sh
+#
+#   Single node (LoRA 1.8B):
+#     YAML_FILE=hy_dense_1_8b_lora_sft.yaml bash train_lf_dense.sh
+#
+#   Multi-node:
+#     Run this script on EACH node with the same IP_LIST.
+#     IP_LIST="10.0.0.1,10.0.0.2" bash train_lf_dense.sh
+# ============================================================================
+
+set -euo pipefail
+
+# -------------------- Network Configuration --------------------
+NET_TYPE="high"
+export NCCL_DEBUG=WARN
+export NCCL_P2P_LEVEL=NVL
+export NCCL_IB_TIMEOUT=24
+export NCCL_NVLS_ENABLE=0
+export NCCL_MPI_PROFILE_PRIMS_ENABLE=0
+export CUDA_DEVICE_MAX_CONNECTIONS=1
+export TORCH_NCCL_HEARTBEAT_TIMEOUT_SEC=3600
+if [[ "${NET_TYPE}" = "low" ]]; then
+    export NCCL_SOCKET_IFNAME=eth1
+    export NCCL_IB_GID_INDEX=3
+    export NCCL_IB_HCA=mlx5_2:1
+    export NCCL_IB_SL=3
+    export NCCL_CHECK_DISABLE=1
+    export NCCL_P2P_DISABLE=0
+    export NCCL_LL_THRESHOLD=16384
+    export NCCL_IB_CUDA_SUPPORT=1
+else
+    export NCCL_IB_GID_INDEX=3
+    export NCCL_IB_SL=3
+    export NCCL_CHECK_DISABLE=1
+    export NCCL_P2P_DISABLE=0
+    export NCCL_IB_DISABLE=0
+    export NCCL_LL_THRESHOLD=16384
+    export NCCL_IB_CUDA_SUPPORT=1
+    export NCCL_SOCKET_IFNAME=bond1
+    export UCX_NET_DEVICES=bond1
+    export NCCL_IB_HCA=mlx5_bond_1,mlx5_bond_5,mlx5_bond_3,mlx5_bond_7,mlx5_bond_4,mlx5_bond_8,mlx5_bond_2,mlx5_bond_6
+    export NCCL_COLLNET_ENABLE=0
+    export SHARP_COLL_ENABLE_SAT=0
+    export NCCL_NET_GDR_LEVEL=2
+    export NCCL_IB_QPS_PER_CONNECTION=4
+    export NCCL_IB_TC=160
+    export NCCL_PXN_DISABLE=1
+fi
+
+# Skip LLaMA Factory version check (we use a newer transformers branch)
+export DISABLE_VERSION_CHECK=1
+
+# -------------------- Node Configuration --------------------
+export HOST_GPU_NUM=8
+# IP list, comma separated. e.g. "10.0.0.1,10.0.0.2" or single node "127.0.0.1"
+export IP_LIST=${IP_LIST:-"127.0.0.1"}
+
+MASTER_PORT=${MASTER_PORT:-29500}
+
+IFS=',' read -ra IP_ARRAY <<< "$IP_LIST"
+NODES=${#IP_ARRAY[@]}
+MASTER_ADDR=${IP_ARRAY[0]}
+
+# -------------------- Paths --------------------
+SCRIPT_DIR="$(cd "$(dirname "${BASH_SOURCE[0]}")" && pwd)"
+YAML_FILE="${YAML_FILE:-${SCRIPT_DIR}/hy_dense_1_8b_full_sft.yaml}"
+ENTRY_SCRIPT="${SCRIPT_DIR}/train_hy_dense.py"
+
+# -------------------- Distributed Environment --------------------
+export MASTER_ADDR="${MASTER_ADDR}"
+export MASTER_PORT="${MASTER_PORT}"
+export NNODES="${NODES}"
+
+if [ ${NODES} -gt 1 ]; then
+    # Determine local node rank by matching local IP against IP_LIST
+    LOCAL_IP=$(hostname -i | awk '{print $1}')
+    NODE_RANK=0
+    for i in "${!IP_ARRAY[@]}"; do
+        if [[ "${IP_ARRAY[$i]}" == "${LOCAL_IP}" ]]; then
+            NODE_RANK=$i
+            break
+        fi
+    done
+    export RANK="${NODE_RANK}"
+else
+    export RANK=0
+fi
+
+echo "============================================"
+echo "  HunYuan Dense LLaMA Factory Training"
+echo "  Nodes: ${NNODES}, Rank: ${RANK}"
+echo "  Master: ${MASTER_ADDR}:${MASTER_PORT}"
+echo "  GPUs per node: ${HOST_GPU_NUM}"
+echo "  Total GPUs: $((NODES * HOST_GPU_NUM))"
+echo "  YAML: ${YAML_FILE}"
+echo "============================================"
+
+# -------------------- Launch --------------------
+# We launch torchrun directly (instead of FORCE_TORCHRUN) so that each
+# worker process runs train_hy_dense.py with all Dense patches applied.
+torchrun \
+    --nnodes "${NNODES}" \
+    --node_rank "${RANK}" \
+    --nproc_per_node "${HOST_GPU_NUM}" \
+    --master_addr "${MASTER_ADDR}" \
+    --master_port "${MASTER_PORT}" \
+    "${ENTRY_SCRIPT}" "${YAML_FILE}"

train/requirements.txt

@@ -0,0 +1,12 @@
+transformers>=5.6.0
+torch>=2.10.0
+torchvision
+torchaudio
+accelerate>=1.11.0
+peft>=0.18.1
+deepspeed>=0.18.7
+sentencepiece
+protobuf
+ninja
+flash-attn
+tensorboard

train/tools/check_converted.py

@@ -0,0 +1,455 @@
+#!/usr/bin/env python3
+"""
+Quick validation script for converted HYV3 outer-format checkpoint.
+
+Checks:
+  1. model.safetensors.index.json structure and completeness
+  2. All expected weight keys exist (dense layer 0, MoE layers 1-79)
+  3. Expert tensor shapes (fused 3D format)
+  4. All referenced shard files exist and are non-empty
+  5. Spot-check: load a few shards and verify tensor shapes/dtypes
+  6. No duplicate or orphan keys
+
+Usage:
+    python check_converted.py <output_dir> [--spot-check N]
+
+Example:
+    python check_converted.py pretrain_base/hf_outer
+    python check_converted.py pretrain_base/hf_outer --spot-check 5
+"""
+
+import argparse
+import json
+import os
+import sys
+import time
+from collections import defaultdict
+
+# ============================================================================
+# Expected key patterns for HYV3 outer format
+# ============================================================================
+
+# Dense layer (layer 0) expected suffixes
+DENSE_SUFFIXES = [
+    "input_layernorm.weight",
+    "post_attention_layernorm.weight",
+    "self_attn.q_proj.weight",
+    "self_attn.k_proj.weight",
+    "self_attn.v_proj.weight",
+    "self_attn.o_proj.weight",
+    "self_attn.q_norm.weight",
+    "self_attn.k_norm.weight",
+    "mlp.gate_proj.weight",
+    "mlp.up_proj.weight",
+    "mlp.down_proj.weight",
+]
+
+# MoE layer (layers 1-79) expected suffixes
+MOE_SUFFIXES = [
+    "input_layernorm.weight",
+    "post_attention_layernorm.weight",
+    "self_attn.q_proj.weight",
+    "self_attn.k_proj.weight",
+    "self_attn.v_proj.weight",
+    "self_attn.o_proj.weight",
+    "self_attn.q_norm.weight",
+    "self_attn.k_norm.weight",
+    # MoE-specific
+    "mlp.gate.weight",
+    "mlp.e_score_correction_bias",
+    "mlp.experts.gate_up_proj",
+    "mlp.experts.down_proj",
+    "mlp.shared_experts.gate_proj.weight",
+    "mlp.shared_experts.up_proj.weight",
+    "mlp.shared_experts.down_proj.weight",
+]
+
+# MTP (Multi-Token Prediction) layer expected suffixes
+# MTP layers share MoE structure but have additional projection/norm keys
+MTP_EXTRA_SUFFIXES = [
+    "eh_proj.weight",
+    "enorm.weight",
+    "final_layernorm.weight",
+    "hnorm.weight",
+]
+
+# Global keys (not per-layer)
+GLOBAL_KEYS = [
+    "model.embed_tokens.weight",
+    "model.norm.weight",
+    "lm_head.weight",
+]
+
+
+def load_config(output_dir):
+    """Load config.json and extract model parameters."""
+    config_path = os.path.join(output_dir, "config.json")
+    if not os.path.exists(config_path):
+        print(f"[ERROR] config.json not found in {output_dir}")
+        return None
+    with open(config_path) as f:
+        return json.load(f)
+
+
+def check_index_json(output_dir):
+    """Check model.safetensors.index.json for structure and completeness."""
+    index_path = os.path.join(output_dir, "model.safetensors.index.json")
+    if not os.path.exists(index_path):
+        print(f"[ERROR] model.safetensors.index.json not found")
+        return None, []
+
+    with open(index_path) as f:
+        index = json.load(f)
+
+    errors = []
+
+    # Check structure
+    if "metadata" not in index:
+        errors.append("Missing 'metadata' in index.json")
+    elif "total_size" not in index["metadata"]:
+        errors.append("Missing 'total_size' in metadata")
+
+    if "weight_map" not in index:
+        errors.append("Missing 'weight_map' in index.json")
+        return index, errors
+
+    weight_map = index["weight_map"]
+    total_size = index.get("metadata", {}).get("total_size", 0)
+
+    print(f"  Index keys     : {len(weight_map)}")
+    print(f"  Total size     : {total_size / 1e9:.2f} GB")
+
+    # Check for empty weight_map
+    if len(weight_map) == 0:
+        errors.append("weight_map is empty")
+
+    return index, errors
+
+
+def check_expected_keys(weight_map, config):
+    """Check that all expected keys exist in the weight_map."""
+    errors = []
+    warnings = []
+
+    num_layers = config.get("num_hidden_layers", 80)
+    first_k_dense = config.get("first_k_dense_replace", 1)
+    num_experts = config.get("num_experts", 192)
+    num_mtp_layers = config.get("num_nextn_predict_layers", 0)
+
+    # Check global keys
+    for key in GLOBAL_KEYS:
+        if key not in weight_map:
+            errors.append(f"Missing global key: {key}")
+
+    # Check per-layer keys (regular layers)
+    missing_by_type = defaultdict(list)
+    for layer_idx in range(num_layers):
+        prefix = f"model.layers.{layer_idx}."
+        if layer_idx < first_k_dense:
+            # Dense layer
+            suffixes = DENSE_SUFFIXES
+        else:
+            # MoE layer
+            suffixes = MOE_SUFFIXES
+
+        for suffix in suffixes:
+            full_key = prefix + suffix
+            if full_key not in weight_map:
+                missing_by_type[suffix].append(layer_idx)
+
+    # Check MTP layers (layer num_layers .. num_layers + num_mtp_layers - 1)
+    mtp_missing_by_type = defaultdict(list)
+    for mtp_idx in range(num_mtp_layers):
+        layer_idx = num_layers + mtp_idx
+        prefix = f"model.layers.{layer_idx}."
+        # MTP layers use MoE structure + extra projection/norm keys
+        mtp_suffixes = MOE_SUFFIXES + MTP_EXTRA_SUFFIXES
+        for suffix in mtp_suffixes:
+            full_key = prefix + suffix
+            if full_key not in weight_map:
+                mtp_missing_by_type[suffix].append(layer_idx)
+
+    for suffix, layers in sorted(mtp_missing_by_type.items()):
+        layer_str = str(layers)
+        errors.append(f"Missing MTP key '{suffix}' in layers: {layer_str}")
+
+    for suffix, layers in sorted(missing_by_type.items()):
+        if len(layers) <= 5:
+            layer_str = str(layers)
+        else:
+            layer_str = f"{layers[:3]}...({len(layers)} total)"
+        errors.append(f"Missing '{suffix}' in layers: {layer_str}")
+
+    # Check for unexpected keys (not matching any known pattern)
+    known_prefixes = set()
+    # Regular layers + MTP layers
+    for layer_idx in range(num_layers + num_mtp_layers):
+        known_prefixes.add(f"model.layers.{layer_idx}.")
+    known_prefixes.add("model.embed_tokens.")
+    known_prefixes.add("model.norm.")
+    known_prefixes.add("lm_head.")
+    # Alternative MTP prefix (some models use this)
+    known_prefixes.add("model.mtp_layers.")
+
+    unexpected = []
+    for key in weight_map:
+        if not any(key.startswith(p) for p in known_prefixes):
+            unexpected.append(key)
+
+    if unexpected:
+        if len(unexpected) <= 5:
+            for k in unexpected:
+                warnings.append(f"Unexpected key: {k}")
+        else:
+            warnings.append(f"{len(unexpected)} unexpected keys found (first 3: {unexpected[:3]})")
+
+    return errors, warnings
+
+
+def check_shard_files(output_dir, weight_map):
+    """Check that all referenced shard files exist and are non-empty."""
+    errors = []
+    warnings = []
+
+    # Get unique shard files
+    shard_files = sorted(set(weight_map.values()))
+    print(f"  Shard files    : {len(shard_files)}")
+
+    missing = []
+    empty = []
+    total_disk_size = 0
+
+    for sf in shard_files:
+        path = os.path.join(output_dir, sf)
+        if not os.path.exists(path):
+            missing.append(sf)
+        else:
+            size = os.path.getsize(path)
+            if size == 0:
+                empty.append(sf)
+            total_disk_size += size
+
+    print(f"  Disk size      : {total_disk_size / 1e9:.2f} GB")
+
+    if missing:
+        errors.append(f"Missing shard files ({len(missing)}): {missing[:5]}")
+    if empty:
+        errors.append(f"Empty shard files ({len(empty)}): {empty[:5]}")
+
+    # Check for orphan shard files (exist on disk but not in index)
+    all_safetensors = set(
+        f for f in os.listdir(output_dir)
+        if f.endswith(".safetensors")
+    )
+    referenced = set(shard_files)
+    orphans = all_safetensors - referenced
+    if orphans:
+        # Distinguish between empty residue files (cross-shard merge artifacts)
+        # and real orphan files with actual data
+        EMPTY_SHARD_THRESHOLD = 128  # bytes; empty safetensors header is ~16 bytes
+        residue_orphans = []
+        real_orphans = []
+        for o in sorted(orphans):
+            sz = os.path.getsize(os.path.join(output_dir, o))
+            if sz <= EMPTY_SHARD_THRESHOLD:
+                residue_orphans.append(o)
+            else:
+                real_orphans.append(o)
+
+        if residue_orphans:
+            warnings.append(
+                f"{len(residue_orphans)} empty residue shard(s) from cross-shard merge "
+                f"(<=128 bytes each, safe to delete)"
+            )
+        if real_orphans:
+            errors.append(
+                f"Orphan shard files with data (not in index): {real_orphans[:5]}"
+            )
+
+    return errors, warnings
+
+
+def check_key_distribution(weight_map):
+    """Check the distribution of keys across shards."""
+    shard_key_count = defaultdict(int)
+    for key, shard in weight_map.items():
+        shard_key_count[shard] += 1
+
+    counts = sorted(shard_key_count.values())
+    print(f"  Keys/shard     : min={counts[0]}, max={counts[-1]}, "
+          f"median={counts[len(counts)//2]}")
+
+    # Check for shards with 0 keys (should not happen if they are in weight_map)
+    zero_shards = [s for s, c in shard_key_count.items() if c == 0]
+    if zero_shards:
+        return [f"Shards with 0 keys: {zero_shards}"]
+    return []
+
+
+def spot_check_shards(output_dir, weight_map, config, num_checks=3):
+    """Spot-check a few shards by loading and verifying tensor shapes."""
+    errors = []
+
+    try:
+        from safetensors import safe_open
+    except ImportError:
+        print("  [SKIP] safetensors not installed, skipping spot-check")
+        return errors
+
+    num_experts = config.get("num_experts", 192)
+    expert_hidden = config.get("expert_hidden_dim", config.get("moe_intermediate_size", 1536))
+    hidden_size = config.get("hidden_size", 4096)
+
+    # Find shards that contain expert tensors (most interesting to check)
+    expert_shards = set()
+    for key, shard in weight_map.items():
+        if "experts.gate_up_proj" in key or "experts.down_proj" in key:
+            expert_shards.add(shard)
+
+    # Pick a few shards to check
+    check_shards = sorted(expert_shards)[:num_checks]
+    if not check_shards:
+        check_shards = sorted(set(weight_map.values()))[:num_checks]
+
+    print(f"\n  Spot-checking {len(check_shards)} shard(s)...")
+
+    for shard_file in check_shards:
+        shard_path = os.path.join(output_dir, shard_file)
+        t0 = time.time()
+
+        try:
+            with safe_open(shard_path, framework="pt", device="cpu") as f:
+                keys_in_shard = list(f.keys())
+                for key in keys_in_shard:
+                    tensor = f.get_tensor(key)
+
+                    # Check expert shapes
+                    if key.endswith("experts.gate_up_proj"):
+                        expected_shape = (num_experts, expert_hidden * 2, hidden_size)
+                        if tuple(tensor.shape) != expected_shape:
+                            errors.append(
+                                f"{shard_file}/{key}: shape {tuple(tensor.shape)} "
+                                f"!= expected {expected_shape}"
+                            )
+
+                    elif key.endswith("experts.down_proj"):
+                        expected_shape = (num_experts, hidden_size, expert_hidden)
+                        if tuple(tensor.shape) != expected_shape:
+                            errors.append(
+                                f"{shard_file}/{key}: shape {tuple(tensor.shape)} "
+                                f"!= expected {expected_shape}"
+                            )
+
+                    # Check for NaN/Inf
+                    if tensor.is_floating_point():
+                        if tensor.isnan().any():
+                            errors.append(f"{shard_file}/{key}: contains NaN values")
+                        if tensor.isinf().any():
+                            errors.append(f"{shard_file}/{key}: contains Inf values")
+
+            elapsed = time.time() - t0
+            print(f"    {shard_file}: {len(keys_in_shard)} keys, OK ({elapsed:.1f}s)")
+
+        except Exception as e:
+            errors.append(f"Failed to load {shard_file}: {e}")
+
+    return errors
+
+
+def main():
+    parser = argparse.ArgumentParser(
+        description="Validate converted HYV3 outer-format checkpoint."
+    )
+    parser.add_argument(
+        "output_dir", type=str,
+        help="Path to the converted outer-format checkpoint directory.",
+    )
+    parser.add_argument(
+        "--spot-check", type=int, default=3, dest="spot_check",
+        help="Number of shards to spot-check by loading tensors (default: 3).",
+    )
+    args = parser.parse_args()
+
+    output_dir = os.path.abspath(args.output_dir)
+    print(f"Validating: {output_dir}\n")
+
+    if not os.path.isdir(output_dir):
+        print(f"[ERROR] Directory not found: {output_dir}")
+        sys.exit(1)
+
+    all_errors = []
+    all_warnings = []
+
+    # 1. Load config
+    print("[1/5] Loading config.json...")
+    config = load_config(output_dir)
+    if config is None:
+        print("[ERROR] Cannot proceed without config.json")
+        sys.exit(1)
+
+    num_layers = config.get("num_hidden_layers", 0)
+    num_experts = config.get("num_experts", 0)
+    first_k_dense = config.get("first_k_dense_replace", 0)
+    num_mtp = config.get("num_nextn_predict_layers", 0)
+    print(f"  Layers         : {num_layers} ({first_k_dense} dense, {num_layers - first_k_dense} MoE)")
+    print(f"  MTP layers     : {num_mtp}")
+    print(f"  Experts/layer  : {num_experts}")
+    print(f"  Hidden size    : {config.get('hidden_size', '?')}")
+    print(f"  Expert hidden  : {config.get('expert_hidden_dim', config.get('moe_intermediate_size', '?'))}")
+
+    # 2. Check index.json
+    print("\n[2/5] Checking model.safetensors.index.json...")
+    index, idx_errors = check_index_json(output_dir)
+    all_errors.extend(idx_errors)
+
+    if index is None or "weight_map" not in index:
+        print("[ERROR] Cannot proceed without valid index.json")
+        sys.exit(1)
+
+    weight_map = index["weight_map"]
+
+    # 3. Check expected keys
+    print("\n[3/5] Checking expected keys...")
+    key_errors, key_warnings = check_expected_keys(weight_map, config)
+    all_errors.extend(key_errors)
+    all_warnings.extend(key_warnings)
+
+    # Also check key distribution
+    dist_errors = check_key_distribution(weight_map)
+    all_errors.extend(dist_errors)
+
+    # 4. Check shard files
+    print("\n[4/5] Checking shard files on disk...")
+    shard_errors, shard_warnings = check_shard_files(output_dir, weight_map)
+    all_errors.extend(shard_errors)
+    all_warnings.extend(shard_warnings)
+
+    # 5. Spot-check
+    if args.spot_check > 0:
+        print(f"\n[5/5] Spot-checking tensors (loading {args.spot_check} shard(s))...")
+        spot_errors = spot_check_shards(output_dir, weight_map, config, args.spot_check)
+        all_errors.extend(spot_errors)
+    else:
+        print("\n[5/5] Spot-check skipped (--spot-check 0)")
+
+    # Summary
+    print(f"\n{'=' * 60}")
+    if all_warnings:
+        print(f"WARNINGS ({len(all_warnings)}):")
+        for w in all_warnings:
+            print(f"  [WARN] {w}")
+
+    if all_errors:
+        print(f"ERRORS ({len(all_errors)}):")
+        for e in all_errors:
+            print(f"  [ERROR] {e}")
+        print(f"\nResult: FAILED ({len(all_errors)} error(s), {len(all_warnings)} warning(s))")
+        sys.exit(1)
+    else:
+        print(f"Result: PASSED (0 errors, {len(all_warnings)} warning(s))")
+        print(f"{'=' * 60}")
+        sys.exit(0)
+
+
+if __name__ == "__main__":
+    main()

train/tools/convert_ckpt_to_outer.py

@@ -0,0 +1,641 @@
+#!/usr/bin/env python3
+"""
+Memory-friendly checkpoint converter: inner -> outer format (v2).
+
+Converts the HYV3 checkpoint from inner format (per-expert keys, old naming)
+to outer format (fused 3D experts, new naming) shard by shard.
+
+Handles the case where a single layer's experts may be split across
+multiple shards (cross-shard experts) by deferring their fusion to a
+post-processing step.
+
+v2 improvements over v1:
+  - Post-processing is shard-centric (each shard read/written only once)
+    instead of prefix-centric (same shard read/written multiple times).
+    This fixes Bus error (core dump) when there are many cross-shard groups.
+  - Explicit memory management with gc.collect() to prevent memory bloat.
+  - Better progress reporting during post-processing.
+
+Supports multi-process parallelism for faster conversion.
+
+Usage:
+    # Default 8 workers
+    python convert_ckpt_to_outer.py \\
+        --input_dir pretrain_base/hf \\
+        --output_dir pretrain_base/hf_outer
+
+    # Custom worker count
+    python convert_ckpt_to_outer.py \\
+        --input_dir pretrain_base/hf \\
+        --output_dir pretrain_base/hf_outer \\
+        --workers 16
+
+The script will:
+  1. Pre-scan index.json to detect cross-shard expert groups
+  2. Convert weights shard-by-shard in parallel (key rename + expert fuse)
+  3. Post-process cross-shard expert groups (merge from multiple shards)
+     - v2: shard-centric approach, each shard read/written only once
+  4. Copy config.json as-is (already in outer format)
+  5. Copy all other files (tokenizer, etc.)
+  6. Rebuild model.safetensors.index.json
+"""
+
+import argparse
+import gc
+import json
+import os
+import re
+import signal
+import shutil
+import sys
+import time
+import traceback
+from collections import OrderedDict, defaultdict
+from multiprocessing import Pool
+
+import torch
+
+try:
+    from safetensors import safe_open
+    from safetensors.torch import save_file
+except ImportError:
+    raise ImportError("Please install safetensors: pip install safetensors")
+
+# ============================================================================
+# Signal handling for Bus error (SIGBUS) and other fatal signals
+# ============================================================================
+
+def _fatal_signal_handler(signum, frame):
+    """Handle fatal signals (SIGBUS, SIGSEGV) by logging before exit.
+
+    These signals cannot be caught by try/except. This handler ensures
+    the error message is written to stderr (captured by nohup redirection)
+    before the process terminates.
+    """
+    sig_name = signal.Signals(signum).name if hasattr(signal, 'Signals') else str(signum)
+    pid = os.getpid()
+    msg = (
+        f"\n[FATAL] Process {pid} received {sig_name} (signal {signum}).\n"
+        f"This typically indicates an out-of-memory condition during mmap I/O.\n"
+        f"Stack trace at time of signal:\n"
+    )
+    sys.stderr.write(msg)
+    traceback.print_stack(frame, file=sys.stderr)
+    sys.stderr.flush()
+    # Re-raise with default handler to get proper exit code
+    signal.signal(signum, signal.SIG_DFL)
+    os.kill(pid, signum)
+
+
+def _install_signal_handlers():
+    """Install handlers for SIGBUS and SIGSEGV in the current process."""
+    for sig in (signal.SIGBUS, signal.SIGSEGV):
+        try:
+            signal.signal(sig, _fatal_signal_handler)
+        except (OSError, ValueError):
+            # Some signals may not be available on all platforms
+            pass
+
+
+def _pool_worker_init():
+    """Initializer for multiprocessing pool workers.
+
+    Installs signal handlers so that Bus errors in worker processes
+    are also logged before the process dies.
+    """
+    _install_signal_handlers()
+
+
+# ============================================================================
+# Key rename mapping (inner -> outer)
+# ============================================================================
+
+_KEY_RENAMES = [
+    ("mlp.router.gate.", "mlp.gate."),
+    ("mlp.expert_bias", "mlp.e_score_correction_bias"),
+    ("mlp.shared_mlp.", "mlp.shared_experts."),
+]
+
+# Regex to match per-expert keys
+_EXPERT_KEY_RE = re.compile(
+    r"^(.*\.mlp\.experts\.)(\d+)\.(gate_proj|up_proj|down_proj)\.weight$"
+)
+
+def rename_key(key: str) -> str:
+    """Rename a single key from inner to outer format."""
+    for old_sub, new_sub in _KEY_RENAMES:
+        if old_sub in key:
+            key = key.replace(old_sub, new_sub)
+            break
+    return key
+
+def scan_cross_shard_experts(index_path: str):
+    """Pre-scan index.json to find expert groups that span multiple shards.
+
+    Returns:
+        cross_shard_prefixes: set of expert prefixes that span multiple shards
+            e.g. {"model.layers.80.mlp.experts."}
+    """
+    with open(index_path) as f:
+        index = json.load(f)
+    wm = index["weight_map"]
+
+    # prefix -> set of shards
+    prefix_shards = defaultdict(set)
+    for key in wm:
+        m = _EXPERT_KEY_RE.match(key)
+        if m:
+            prefix = m.group(1)
+            prefix_shards[prefix].add(wm[key])
+
+    cross_shard_prefixes = set()
+    for prefix, shards in prefix_shards.items():
+        if len(shards) > 1:
+            cross_shard_prefixes.add(prefix)
+
+    return cross_shard_prefixes
+
+def convert_shard(shard_path: str, cross_shard_prefixes: set = None):
+    """Load a single shard, rename keys, and fuse experts.
+
+    For expert groups in cross_shard_prefixes, the per-expert keys are
+    kept as-is (just renamed) and returned separately as deferred items,
+    to be merged later in a post-processing step.
+
+    Returns:
+        result: OrderedDict of converted tensors (ready to save)
+        deferred_expert_keys: list of original expert keys that were deferred
+            (these are kept in result with their original per-expert naming
+             but with the outer rename applied, to be post-processed later)
+    """
+    if cross_shard_prefixes is None:
+        cross_shard_prefixes = set()
+
+    tensors = OrderedDict()
+    with safe_open(shard_path, framework="pt", device="cpu") as f:
+        for key in f.keys():
+            tensors[key] = f.get_tensor(key)
+
+    # Separate expert keys from non-expert keys
+    expert_groups = {}  # prefix -> {expert_idx -> {proj_name -> tensor}}
+    deferred_expert_keys = []  # keys that belong to cross-shard experts
+    result = OrderedDict()
+
+    for key, tensor in tensors.items():
+        m = _EXPERT_KEY_RE.match(key)
+        if m:
+            prefix = m.group(1)
+            expert_idx = int(m.group(2))
+            proj_name = m.group(3)
+
+            if prefix in cross_shard_prefixes:
+                # Defer: keep the key as-is (with rename) for post-processing
+                new_key = rename_key(key)
+                result[new_key] = tensor
+                deferred_expert_keys.append(new_key)
+            else:
+                # Normal: collect for fusion within this shard
+                if prefix not in expert_groups:
+                    expert_groups[prefix] = {}
+                if expert_idx not in expert_groups[prefix]:
+                    expert_groups[prefix][expert_idx] = {}
+                expert_groups[prefix][expert_idx][proj_name] = tensor
+        else:
+            # Non-expert key: just rename
+            new_key = rename_key(key)
+            result[new_key] = tensor
+
+    # Fuse expert weights for each non-cross-shard layer prefix
+    for prefix in sorted(expert_groups.keys()):
+        experts = expert_groups[prefix]
+        num_experts = max(experts.keys()) + 1
+
+        gate_up_list = []
+        down_list = []
+        for i in range(num_experts):
+            if i not in experts:
+                raise ValueError(
+                    f"Missing expert {i} in {prefix}. "
+                    f"Found: {sorted(experts.keys())}"
+                )
+            exp = experts[i]
+            gate_up = torch.cat([exp["gate_proj"], exp["up_proj"]], dim=0)
+            gate_up_list.append(gate_up)
+            down_list.append(exp["down_proj"])
+
+        fused_gate_up = torch.stack(gate_up_list, dim=0)
+        fused_down = torch.stack(down_list, dim=0)
+
+        for exp in experts.values():
+            exp.clear()
+        gate_up_list.clear()
+        down_list.clear()
+
+        result[f"{prefix}gate_up_proj"] = fused_gate_up
+        result[f"{prefix}down_proj"] = fused_down
+
+    return result, deferred_expert_keys
+
+def _process_one_shard(args_tuple):
+    """Worker function: convert a single shard and save to output dir.
+
+    Args:
+        args_tuple: (idx, num_shards, shard_file, input_dir, output_dir, cross_shard_prefixes)
+
+    Returns:
+        (shard_file, key_list, shard_size, elapsed, deferred_keys)
+    """
+    idx, num_shards, shard_file, input_dir, output_dir, cross_shard_prefixes = args_tuple
+    shard_path = os.path.join(input_dir, shard_file)
+    t0 = time.time()
+
+    converted, deferred_keys = convert_shard(shard_path, cross_shard_prefixes)
+
+    shard_size = sum(t.numel() * t.element_size() for t in converted.values())
+
+    out_shard_path = os.path.join(output_dir, shard_file)
+    save_file(converted, out_shard_path)
+
+    elapsed = time.time() - t0
+    num_keys = len(converted)
+    key_list = list(converted.keys())
+
+    del converted
+
+    deferred_info = ""
+    if deferred_keys:
+        deferred_info = f", Deferred={len(deferred_keys)}"
+
+    print(
+        f"  [{idx + 1}/{num_shards}] {shard_file}: "
+        f"Keys={num_keys}, Size={shard_size / 1e9:.2f} GB, "
+        f"Time={elapsed:.1f}s{deferred_info}",
+        flush=True,
+    )
+
+    return shard_file, key_list, shard_size, elapsed, deferred_keys
+
+
+def post_process_cross_shard_experts(output_dir, cross_shard_prefixes, all_deferred):
+    """Merge cross-shard expert groups (v2: shard-centric approach).
+
+    Instead of iterating per-prefix (which causes the same shard to be
+    loaded/saved multiple times), this v2 approach:
+      1. Builds a mapping of which prefixes each shard is involved in
+      2. Collects all expert tensors from all involved shards in ONE pass
+      3. Fuses all prefixes
+      4. Writes each shard only ONCE with all its updates applied
+
+    This avoids the Bus error (core dump) caused by repeated mmap of
+    large files and memory bloat.
+
+    Args:
+        output_dir: path to output directory
+        cross_shard_prefixes: set of expert prefixes that span multiple shards
+        all_deferred: dict of {shard_file: [deferred_key, ...]}
+
+    Returns:
+        updated_shards: dict of {shard_file: (key_list, shard_size)} for updated shards
+    """
+    if not cross_shard_prefixes:
+        return {}
+
+    print(f"\n  Post-processing {len(cross_shard_prefixes)} cross-shard expert group(s)...",
+          flush=True)
+
+    # ----------------------------------------------------------------
+    # Step 1: Build mappings
+    # ----------------------------------------------------------------
+    # prefix -> ordered list of shards that contain its experts
+    prefix_to_shards = defaultdict(set)
+    # shard -> set of prefixes it is involved in
+    shard_to_prefixes = defaultdict(set)
+
+    for shard_file, deferred_keys in all_deferred.items():
+        for key in deferred_keys:
+            m = _EXPERT_KEY_RE.match(key)
+            if m:
+                prefix = m.group(1)
+                if prefix in cross_shard_prefixes:
+                    prefix_to_shards[prefix].add(shard_file)
+                    shard_to_prefixes[shard_file].add(prefix)
+
+    # For each prefix, decide which shard will hold the fused result
+    # (use the first shard alphabetically)
+    prefix_to_target_shard = {}
+    for prefix in sorted(prefix_to_shards.keys()):
+        target = sorted(prefix_to_shards[prefix])[0]
+        prefix_to_target_shard[prefix] = target
+
+    # All shards that need to be updated
+    all_involved_shards = set()
+    for shards in prefix_to_shards.values():
+        all_involved_shards.update(shards)
+
+    print(f"    Involved shards: {len(all_involved_shards)}", flush=True)
+    print(f"    Expert groups: {len(prefix_to_shards)}", flush=True)
+
+    # ----------------------------------------------------------------
+    # Step 2: Collect all expert tensors from all involved shards
+    #         (one pass per shard)
+    # ----------------------------------------------------------------
+    # prefix -> {expert_idx -> {proj_name -> tensor}}
+    all_expert_data = defaultdict(dict)
+    # shard -> OrderedDict of non-expert keys (to be re-saved)
+    shard_non_expert = {}
+
+    sorted_involved = sorted(all_involved_shards)
+    for si, shard_file in enumerate(sorted_involved):
+        shard_path = os.path.join(output_dir, shard_file)
+        prefixes_in_shard = shard_to_prefixes[shard_file]
+
+        print(f"    [{si+1}/{len(sorted_involved)}] Reading {shard_file} "
+              f"({len(prefixes_in_shard)} prefix(es))...", flush=True)
+
+        non_expert = OrderedDict()
+        with safe_open(shard_path, framework="pt", device="cpu") as f:
+            for key in f.keys():
+                m = _EXPERT_KEY_RE.match(key)
+                if m and m.group(1) in prefixes_in_shard:
+                    # This is a deferred expert key
+                    prefix = m.group(1)
+                    expert_idx = int(m.group(2))
+                    proj_name = m.group(3)
+                    if expert_idx not in all_expert_data[prefix]:
+                        all_expert_data[prefix][expert_idx] = {}
+                    all_expert_data[prefix][expert_idx][proj_name] = f.get_tensor(key)
+                else:
+                    # Non-expert key: keep as-is
+                    non_expert[key] = f.get_tensor(key)
+
+        shard_non_expert[shard_file] = non_expert
+        gc.collect()
+
+    # ----------------------------------------------------------------
+    # Step 3: Fuse all expert groups
+    # ----------------------------------------------------------------
+    # prefix -> {"gate_up_proj": tensor, "down_proj": tensor}
+    fused_results = {}
+
+    for pi, prefix in enumerate(sorted(all_expert_data.keys())):
+        expert_data = all_expert_data[prefix]
+        num_experts = max(expert_data.keys()) + 1
+
+        print(f"    Fusing {prefix} ({num_experts} experts)...", flush=True)
+
+        gate_up_list = []
+        down_list = []
+        for i in range(num_experts):
+            if i not in expert_data:
+                raise ValueError(
+                    f"Missing expert {i} in {prefix} after cross-shard merge. "
+                    f"Found: {sorted(expert_data.keys())}"
+                )
+            exp = expert_data[i]
+            if "gate_proj" not in exp or "up_proj" not in exp:
+                raise ValueError(
+                    f"Expert {i} in {prefix} missing gate_proj/up_proj. "
+                    f"Has: {sorted(exp.keys())}"
+                )
+            if "down_proj" not in exp:
+                raise ValueError(
+                    f"Expert {i} in {prefix} missing down_proj. "
+                    f"Has: {sorted(exp.keys())}"
+                )
+            gate_up = torch.cat([exp["gate_proj"], exp["up_proj"]], dim=0)
+            gate_up_list.append(gate_up)
+            down_list.append(exp["down_proj"])
+
+        fused_gate_up = torch.stack(gate_up_list, dim=0)
+        fused_down = torch.stack(down_list, dim=0)
+
+        fused_results[prefix] = {
+            "gate_up_proj": fused_gate_up,
+            "down_proj": fused_down,
+        }
+
+        # Free per-expert data for this prefix
+        del gate_up_list, down_list
+        for exp in expert_data.values():
+            exp.clear()
+        del all_expert_data[prefix]
+        gc.collect()
+
+    del all_expert_data
+    gc.collect()
+
+    # ----------------------------------------------------------------
+    # Step 4: Write each involved shard ONCE with all updates applied
+    # ----------------------------------------------------------------
+    updated_shards = {}
+
+    for si, shard_file in enumerate(sorted_involved):
+        shard_path = os.path.join(output_dir, shard_file)
+        non_expert = shard_non_expert[shard_file]
+
+        # Add fused tensors for prefixes that target this shard
+        fused_added = []
+        for prefix, target_shard in prefix_to_target_shard.items():
+            if target_shard == shard_file and prefix in fused_results:
+                non_expert[f"{prefix}gate_up_proj"] = fused_results[prefix]["gate_up_proj"]
+                non_expert[f"{prefix}down_proj"] = fused_results[prefix]["down_proj"]
+                fused_added.append(prefix)
+
+        save_file(non_expert, shard_path)
+        shard_size = sum(t.numel() * t.element_size() for t in non_expert.values())
+        updated_shards[shard_file] = (list(non_expert.keys()), shard_size)
+
+        fused_info = ""
+        if fused_added:
+            fused_info = f", Fused {len(fused_added)} group(s)"
+
+        print(f"    [{si+1}/{len(sorted_involved)}] Wrote {shard_file}: "
+              f"{len(non_expert)} keys, {shard_size / 1e9:.2f} GB{fused_info}",
+              flush=True)
+
+        # Free memory for this shard
+        del shard_non_expert[shard_file]
+        for prefix in fused_added:
+            del fused_results[prefix]
+        del non_expert
+        gc.collect()
+
+    return updated_shards
+
+
+def main():
+    parser = argparse.ArgumentParser(
+        description="Convert HYV3 checkpoint from inner to outer format (v2, shard-centric post-processing)."
+    )
+    parser.add_argument(
+        "--input_dir", type=str, required=True,
+        help="Path to the inner-format checkpoint directory.",
+    )
+    parser.add_argument(
+        "--output_dir", type=str, required=True,
+        help="Path to the output outer-format checkpoint directory.",
+    )
+    parser.add_argument(
+        "--workers", type=int, default=8,
+        help="Number of parallel worker processes (default: 8).",
+    )
+    args = parser.parse_args()
+
+    input_dir = os.path.abspath(args.input_dir)
+    output_dir = os.path.abspath(args.output_dir)
+    num_workers = args.workers
+
+    if not os.path.isdir(input_dir):
+        raise FileNotFoundError(f"Input directory not found: {input_dir}")
+
+    os.makedirs(output_dir, exist_ok=True)
+
+    # Pre-scan for cross-shard expert groups
+    index_path = os.path.join(input_dir, "model.safetensors.index.json")
+    cross_shard_prefixes = set()
+    if os.path.exists(index_path):
+        cross_shard_prefixes = scan_cross_shard_experts(index_path)
+        if cross_shard_prefixes:
+            print(f"Detected {len(cross_shard_prefixes)} cross-shard expert group(s):")
+            for p in sorted(cross_shard_prefixes):
+                print(f"  - {p}")
+            print()
+
+    # Get all safetensors files
+    shard_files = sorted(
+        f for f in os.listdir(input_dir) if f.endswith(".safetensors")
+    )
+    if not shard_files:
+        raise FileNotFoundError(f"No .safetensors files found in {input_dir}")
+
+    # Skip already-converted shards (for resumability)
+    # NOTE: if there are cross-shard experts, we cannot skip shards that
+    # contain deferred keys (they need post-processing). For simplicity,
+    # when cross-shard experts exist, we re-process all shards.
+    remaining = []
+    skipped = []
+    if cross_shard_prefixes:
+        # Re-process all shards when cross-shard experts exist
+        remaining = list(shard_files)
+    else:
+        for sf in shard_files:
+            out_path = os.path.join(output_dir, sf)
+            if os.path.exists(out_path) and os.path.getsize(out_path) > 0:
+                skipped.append(sf)
+            else:
+                remaining.append(sf)
+
+    num_shards = len(shard_files)
+    num_workers = min(num_workers, len(remaining)) if remaining else 1
+
+    print(f"=" * 60)
+    print(f"HYV3 Checkpoint Converter (inner -> outer, v2)")
+    print(f"  Input  : {input_dir}")
+    print(f"  Output : {output_dir}")
+    print(f"  Shards : {num_shards} total, {len(skipped)} already done, {len(remaining)} to process")
+    print(f"  Workers: {num_workers}")
+    if cross_shard_prefixes:
+        print(f"  Cross-shard experts: {len(cross_shard_prefixes)} group(s) (will post-process)")
+    print(f"=" * 60)
+
+    t_start = time.time()
+
+    # Build task list for remaining shards
+    tasks = [
+        (i, len(remaining), sf, input_dir, output_dir, cross_shard_prefixes)
+        for i, sf in enumerate(remaining)
+    ]
+
+    # Process in parallel
+    results = []
+    if tasks:
+        with Pool(processes=num_workers, initializer=_pool_worker_init) as pool:
+            results = pool.map(_process_one_shard, tasks)
+
+    # Collect deferred keys info
+    all_deferred = {}  # shard_file -> [deferred_keys]
+    for shard_file, key_list, shard_size, elapsed, deferred_keys in results:
+        if deferred_keys:
+            all_deferred[shard_file] = deferred_keys
+
+    # Post-process cross-shard expert groups (v2: shard-centric)
+    updated_shards = {}
+    if cross_shard_prefixes and all_deferred:
+        updated_shards = post_process_cross_shard_experts(
+            output_dir, cross_shard_prefixes, all_deferred
+        )
+
+    # Build weight_map and total_size
+    weight_map = OrderedDict()
+    total_size = 0
+
+    # For skipped shards, read their keys from the output files
+    for sf in skipped:
+        out_path = os.path.join(output_dir, sf)
+        with safe_open(out_path, framework="pt", device="cpu") as f:
+            keys = list(f.keys())
+            for key in keys:
+                weight_map[key] = sf
+                t = f.get_tensor(key)
+                total_size += t.numel() * t.element_size()
+
+    # Collect results from newly converted shards
+    for shard_file, key_list, shard_size, elapsed, deferred_keys in results:
+        if shard_file in updated_shards:
+            # This shard was updated by post-processing
+            updated_key_list, updated_size = updated_shards[shard_file]
+            for key in updated_key_list:
+                weight_map[key] = shard_file
+            total_size += updated_size
+        else:
+            for key in key_list:
+                weight_map[key] = shard_file
+            total_size += shard_size
+
+    # Build and save index
+    sorted_weight_map = OrderedDict(sorted(weight_map.items()))
+    index = {
+        "metadata": {"total_size": total_size},
+        "weight_map": sorted_weight_map,
+    }
+    index_path_out = os.path.join(output_dir, "model.safetensors.index.json")
+    with open(index_path_out, "w") as f:
+        json.dump(index, f, indent=2)
+        f.write("\n")
+    print(f"\nSaved {index_path_out}")
+
+    # Copy non-safetensors files (config, tokenizer, etc.)
+    skip_suffixes = {".safetensors"}
+    skip_names = {"model.safetensors.index.json"}
+    copied = []
+    for fname in os.listdir(input_dir):
+        if fname in skip_names:
+            continue
+        if any(fname.endswith(s) for s in skip_suffixes):
+            continue
+        src = os.path.join(input_dir, fname)
+        dst = os.path.join(output_dir, fname)
+        if os.path.isfile(src):
+            shutil.copy2(src, dst)
+            copied.append(fname)
+        elif os.path.isdir(src):
+            if os.path.exists(dst):
+                shutil.rmtree(dst)
+            shutil.copytree(src, dst)
+            copied.append(fname + "/")
+
+    if copied:
+        print(f"\nCopied files: {', '.join(copied)}")
+
+    t_total = time.time() - t_start
+    print(f"\n{'=' * 60}")
+    print(f"Conversion complete!")
+    print(f"  Total keys : {len(weight_map)}")
+    print(f"  Total size : {total_size / 1e9:.2f} GB")
+    print(f"  Total time : {t_total:.1f}s ({t_total / 60:.1f} min)")
+    print(f"  Output dir : {output_dir}")
+    print(f"{'=' * 60}")
+
+if __name__ == "__main__":
+    _install_signal_handlers()
+    main()