Upload folder using huggingface_hub

.gitattributes

@@ -33,3 +33,5 @@ saved_model/**/* filter=lfs diff=lfs merge=lfs -text
 *.zip filter=lfs diff=lfs merge=lfs -text
 *.zst filter=lfs diff=lfs merge=lfs -text
 *tfevents* filter=lfs diff=lfs merge=lfs -text
+assets/benchmark-appendix.png filter=lfs diff=lfs merge=lfs -text
+assets/benchmark.png filter=lfs diff=lfs merge=lfs -text

LICENSE

@@ -0,0 +1,210 @@
+Tencent is pleased to support the open source community by making Tencent Hy3 available.
+
+Copyright (C) 2026 Tencent.  All rights reserved.
+
+Tencent Hy3 is licensed under the Apache-2.0.
+
+
+Terms of the Apache-2.0:
+--------------------------------------------------------------------
+Apache License
+                           Version 2.0, January 2004
+                        http://www.apache.org/licenses/
+
+   TERMS AND CONDITIONS FOR USE, REPRODUCTION, AND DISTRIBUTION
+
+   1. Definitions.
+
+      "License" shall mean the terms and conditions for use, reproduction,
+      and distribution as defined by Sections 1 through 9 of this document.
+
+      "Licensor" shall mean the copyright owner or entity authorized by
+      the copyright owner that is granting the License.
+
+      "Legal Entity" shall mean the union of the acting entity and all
+      other entities that control, are controlled by, or are under common
+      control with that entity. For the purposes of this definition,
+      "control" means (i) the power, direct or indirect, to cause the
+      direction or management of such entity, whether by contract or
+      otherwise, or (ii) ownership of fifty percent (50%) or more of the
+      outstanding shares, or (iii) beneficial ownership of such entity.
+
+      "You" (or "Your") shall mean an individual or Legal Entity
+      exercising permissions granted by this License.
+
+      "Source" form shall mean the preferred form for making modifications,
+      including but not limited to software source code, documentation
+      source, and configuration files.
+
+      "Object" form shall mean any form resulting from mechanical
+      transformation or translation of a Source form, including but
+      not limited to compiled object code, generated documentation,
+      and conversions to other media types.
+
+      "Work" shall mean the work of authorship, whether in Source or
+      Object form, made available under the License, as indicated by a
+      copyright notice that is included in or attached to the work
+      (an example is provided in the Appendix below).
+
+      "Derivative Works" shall mean any work, whether in Source or Object
+      form, that is based on (or derived from) the Work and for which the
+      editorial revisions, annotations, elaborations, or other modifications
+      represent, as a whole, an original work of authorship. For the purposes
+      of this License, Derivative Works shall not include works that remain
+      separable from, or merely link (or bind by name) to the interfaces of,
+      the Work and Derivative Works thereof.
+
+      "Contribution" shall mean any work of authorship, including
+      the original version of the Work and any modifications or additions
+      to that Work or Derivative Works thereof, that is intentionally
+      submitted to Licensor for inclusion in the Work by the copyright owner
+      or by an individual or Legal Entity authorized to submit on behalf of
+      the copyright owner. For the purposes of this definition, "submitted"
+      means any form of electronic, verbal, or written communication sent
+      to the Licensor or its representatives, including but not limited to
+      communication on electronic mailing lists, source code control systems,
+      and issue tracking systems that are managed by, or on behalf of, the
+      Licensor for the purpose of discussing and improving the Work, but
+      excluding communication that is conspicuously marked or otherwise
+      designated in writing by the copyright owner as "Not a Contribution."
+
+      "Contributor" shall mean Licensor and any individual or Legal Entity
+      on behalf of whom a Contribution has been received by Licensor and
+      subsequently incorporated within the Work.
+
+   2. Grant of Copyright License. Subject to the terms and conditions of
+      this License, each Contributor hereby grants to You a perpetual,
+      worldwide, non-exclusive, no-charge, royalty-free, irrevocable
+      copyright license to reproduce, prepare Derivative Works of,
+      publicly display, publicly perform, sublicense, and distribute the
+      Work and such Derivative Works in Source or Object form.
+
+   3. Grant of Patent License. Subject to the terms and conditions of
+      this License, each Contributor hereby grants to You a perpetual,
+      worldwide, non-exclusive, no-charge, royalty-free, irrevocable
+      (except as stated in this section) patent license to make, have made,
+      use, offer to sell, sell, import, and otherwise transfer the Work,
+      where such license applies only to those patent claims licensable
+      by such Contributor that are necessarily infringed by their
+      Contribution(s) alone or by combination of their Contribution(s)
+      with the Work to which such Contribution(s) was submitted. If You
+      institute patent litigation against any entity (including a
+      cross-claim or counterclaim in a lawsuit) alleging that the Work
+      or a Contribution incorporated within the Work constitutes direct
+      or contributory patent infringement, then any patent licenses
+      granted to You under this License for that Work shall terminate
+      as of the date such litigation is filed.
+
+   4. Redistribution. You may reproduce and distribute copies of the
+      Work or Derivative Works thereof in any medium, with or without
+      modifications, and in Source or Object form, provided that You
+      meet the following conditions:
+
+      (a) You must give any other recipients of the Work or
+          Derivative Works a copy of this License; and
+
+      (b) You must cause any modified files to carry prominent notices
+          stating that You changed the files; and
+
+      (c) You must retain, in the Source form of any Derivative Works
+          that You distribute, all copyright, patent, trademark, and
+          attribution notices from the Source form of the Work,
+          excluding those notices that do not pertain to any part of
+          the Derivative Works; and
+
+      (d) If the Work includes a "NOTICE" text file as part of its
+          distribution, then any Derivative Works that You distribute must
+          include a readable copy of the attribution notices contained
+          within such NOTICE file, excluding those notices that do not
+          pertain to any part of the Derivative Works, in at least one
+          of the following places: within a NOTICE text file distributed
+          as part of the Derivative Works; within the Source form or
+          documentation, if provided along with the Derivative Works; or,
+          within a display generated by the Derivative Works, if and
+          wherever such third-party notices normally appear. The contents
+          of the NOTICE file are for informational purposes only and
+          do not modify the License. You may add Your own attribution
+          notices within Derivative Works that You distribute, alongside
+          or as an addendum to the NOTICE text from the Work, provided
+          that such additional attribution notices cannot be construed
+          as modifying the License.
+
+      You may add Your own copyright statement to Your modifications and
+      may provide additional or different license terms and conditions
+      for use, reproduction, or distribution of Your modifications, or
+      for any such Derivative Works as a whole, provided Your use,
+      reproduction, and distribution of the Work otherwise complies with
+      the conditions stated in this License.
+
+   5. Submission of Contributions. Unless You explicitly state otherwise,
+      any Contribution intentionally submitted for inclusion in the Work
+      by You to the Licensor shall be under the terms and conditions of
+      this License, without any additional terms or conditions.
+      Notwithstanding the above, nothing herein shall supersede or modify
+      the terms of any separate license agreement you may have executed
+      with Licensor regarding such Contributions.
+
+   6. Trademarks. This License does not grant permission to use the trade
+      names, trademarks, service marks, or product names of the Licensor,
+      except as required for reasonable and customary use in describing the
+      origin of the Work and reproducing the content of the NOTICE file.
+
+   7. Disclaimer of Warranty. Unless required by applicable law or
+      agreed to in writing, Licensor provides the Work (and each
+      Contributor provides its Contributions) on an "AS IS" BASIS,
+      WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or
+      implied, including, without limitation, any warranties or conditions
+      of TITLE, NON-INFRINGEMENT, MERCHANTABILITY, or FITNESS FOR A
+      PARTICULAR PURPOSE. You are solely responsible for determining the
+      appropriateness of using or redistributing the Work and assume any
+      risks associated with Your exercise of permissions under this License.
+
+   8. Limitation of Liability. In no event and under no legal theory,
+      whether in tort (including negligence), contract, or otherwise,
+      unless required by applicable law (such as deliberate and grossly
+      negligent acts) or agreed to in writing, shall any Contributor be
+      liable to You for damages, including any direct, indirect, special,
+      incidental, or consequential damages of any character arising as a
+      result of this License or out of the use or inability to use the
+      Work (including but not limited to damages for loss of goodwill,
+      work stoppage, computer failure or malfunction, or any and all
+      other commercial damages or losses), even if such Contributor
+      has been advised of the possibility of such damages.
+
+   9. Accepting Warranty or Additional Liability. While redistributing
+      the Work or Derivative Works thereof, You may choose to offer,
+      and charge a fee for, acceptance of support, warranty, indemnity,
+      or other liability obligations and/or rights consistent with this
+      License. However, in accepting such obligations, You may act only
+      on Your own behalf and on Your sole responsibility, not on behalf
+      of any other Contributor, and only if You agree to indemnify,
+      defend, and hold each Contributor harmless for any liability
+      incurred by, or claims asserted against, such Contributor by reason
+      of your accepting any such warranty or additional liability.
+
+   END OF TERMS AND CONDITIONS
+
+   APPENDIX: How to apply the Apache License to your work.
+
+      To apply the Apache License to your work, attach the following
+      boilerplate notice, with the fields enclosed by brackets "[]"
+      replaced with your own identifying information. (Don't include
+      the brackets!)  The text should be enclosed in the appropriate
+      comment syntax for the file format. We also recommend that a
+      file or class name and description of purpose be included on the
+      same "printed page" as the copyright notice for easier
+      identification within third-party archives.
+
+   Copyright [yyyy] [name of copyright owner]
+
+   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.

finetune/README.md

@@ -0,0 +1,416 @@
+<p align="left">
+    <a href="README_CN.md">中文</a> ｜ English
+</p>
+
+# Model Training
+
+Hy3 provides processes related to model training. This section details how to process training data for model training purposes.
+
+## Training Data Format and Processing
+
+**Hy3 supports both "slow thinking" and "fast thinking" modes. By default, the model outputs in slow thinking mode. If you wish the model to use fast thinking, you can control it via the `reasoning_effort` parameter (options: `high`, `low`, `no_think`).**
+
+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.
+
+```python
+# Fast thinking pattern (no_think)
+{"reasoning_effort": "no_think", "messages": [{"content": "You are a helpful assistant.\nThe current time is 2026-01-01 13:26:12 Thursday", "role": "system"}, {"content": "1+1=?", "role": "user"}, {"role": "assistant", "content": "1+1=2"}]}
+
+# Slow thinking pattern (high)
+{"reasoning_effort": "high", "messages": [{"content": "You are a helpful assistant.\nThe current time is 2026-01-01 13:26:12 Thursday", "role": "system"}, {"content": "1+1=?", "role": "user"}, {"role": "assistant", "content": "1+1=2", "reasoning_content": "The user is asking for the result of 1 + 1. In basic decimal arithmetic, 1 + 1 equals 2."}]}
+
+from transformers import AutoTokenizer
+tokenizer = AutoTokenizer.from_pretrained("./models", use_fast=False, trust_remote_code=True)
+ids = tokenizer.apply_chat_template(messages, is_training=True)
+```
+
+## Checkpoint Format Conversion
+
+The original Hy3 checkpoint stores each expert's weights independently. Before training, it is recommended to convert it to the HuggingFace-compatible format (fusing multiple experts per layer into 3D tensors with unified key naming) to improve loading and training speed. You can also train directly with the original format without conversion, but the loading speed will be slower. We provide a conversion script `convert_ckpt_to_outer.py` and a validation script `check_converted.py`, located in the `train/tools` directory.
+
+### Conversion
+
+```sh
+python convert_ckpt_to_outer.py \
+    --input_dir <original_checkpoint_dir> \
+    --output_dir <output_dir> \
+    --workers 8
+```
+
+**Parameters:**
+
+- `--input_dir`: Path to the original checkpoint directory (required)
+- `--output_dir`: Path to the converted checkpoint output directory (required)
+- `--workers`: Number of parallel worker processes, default is 8 (optional)
+
+The conversion script performs the following steps:
+1. Pre-scans `model.safetensors.index.json` to detect cross-shard expert groups
+2. Converts weights shard-by-shard in parallel (key renaming + expert fusion)
+3. Post-processes cross-shard expert groups (merges data from multiple shards)
+4. Copies `config.json`, tokenizer, and other files
+5. Rebuilds `model.safetensors.index.json`
+
+### Validation
+
+After conversion, it is recommended to validate the result using the validation script:
+
+```sh
+python check_converted.py <converted_checkpoint_dir> --spot-check 3
+```
+
+**Parameters:**
+
+- First argument: Path to the converted checkpoint directory (required)
+- `--spot-check`: Number of shard files to spot-check by loading tensors and verifying shape, dtype, NaN/Inf, etc. Default is 3 (optional)
+
+The validation script checks the following:
+1. Completeness of `config.json`
+2. Whether all expected keys in `model.safetensors.index.json` are present (including regular layers and MTP layers)
+3. Whether all referenced shard files exist and are non-empty
+4. Spot-checks tensor shape, dtype, and NaN/Inf in selected shard files
+5. Detects orphan empty shard files (cross-shard merge residues, safe to delete)
+
+## Quick Start
+
+You can quickly get started by following the instructions in the Quick Start Guide.
+
+## Model Training
+
+### Hardware Requirements
+
+Based on testing, when `make_moe_param_leaf_module` and `zero3+offload` are disabled and `max_seq_length` is set to 4096:
+
+- **LoRA Fine-tuning**: Requires at least a single machine with 8 GPUs (at least 80GB memory each).
+- **Full Fine-tuning**: Requires at least 4 machines with 32 GPUs (at least 80GB memory 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 three 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
+- **ms-swift Training**: Located in the `train/ms_swift_support` directory
+
+#### DeepSpeed Native Training
+
+Reference: [HuggingFace Transformers Trainer](https://huggingface.co/docs/transformers/main/en/main_classes/trainer)
+
+##### Single-Machine Training
+
+In the `train/deepspeed_support` directory, execute:
+
+```sh
+pip install -r requirements.txt
+bash train.sh
+```
+
+##### 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 add the following configuration at the beginning of `train.sh`:
+
+```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 `bash train.sh` 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 HF pre-trained model weights to load, otherwise loading will fail.
+- `--tokenizer_name_or_path`: Path to the tokenizer folder, otherwise loading will fail.
+- `--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.
+
+##### Configuration Files
+
+We provide sample LLaMA-Factory training configuration files: `hy_v3_lora_sft.yaml` and `hy_v3_full_sft.yaml`, corresponding to LoRA training and full fine-tuning respectively.
+
+Key parameters in the configuration files are as follows:
+
+**Model:**
+
+- `model_name_or_path`: Path to the Hy3 HF format pre-trained model weights
+- `trust_remote_code`: Whether to trust remote code; Hy3 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; Hy3 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 `train_lf.sh`:
+
+```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 `bash train_lf.sh` in the `train/llama_factory_support/` directory.
+
+#### ms-swift Training
+
+If you are familiar with ms-swift, you may use it for fine-tuning. All scripts, code, and configuration files are archived in the `train/ms_swift_support` directory. Unless otherwise specified, all files mentioned below are located in this directory.
+
+##### Installation
+
+You can install ms-swift via pip:
+
+```sh
+pip install ms-swift==4.2.2
+```
+
+Or install from source: https://github.com/modelscope/ms-swift
+
+##### Training Scripts and Configuration Files
+
+| Training Method | Configuration File | Launch Script |
+|----------------|-------------------|---------------|
+| Full Fine-tuning | `hy_v3_full_sft.yaml` | `bash sft_train.sh` |
+| LoRA Fine-tuning | `hy_v3_lora_sft.yaml` | `bash sft_train.sh` |
+
+##### About the eos_token_id Patch
+
+The `hy_v3_swift_patches.py` file in the directory fixes an issue with the eos token in ms-swift's default template. The default template uses the `<｜hy_eos｜>` string as `chat_sep` and `suffix`, which gets tokenized into multiple token IDs, causing `model.generate()` to fail to stop correctly during inference.
+
+The patch re-registers the template using the `[['eos_token_id']]` syntax, allowing ms-swift to dynamically resolve `tokenizer.eos_token_id` at runtime and generate the correct single token.
+
+The launch script automatically loads this patch via `--custom_register_path hy_v3_swift_patches.py`, requiring no additional action.
+
+##### Key Parameters
+
+Key parameters in the configuration files are as follows:
+
+**Model:**
+
+- `model`: Model path, can be a HuggingFace Hub ID or a local path
+- `model_type`: Model type, set to `hy_v3`
+- `template`: Chat template, set to `hy_v3`
+- `torch_dtype`: Data type, `bfloat16` is recommended
+- `attn_impl`: Attention implementation, `flash_attn` is recommended
+
+**Training Method:**
+
+- `tuner_type`: Fine-tuning type; set to `full` for full fine-tuning, `lora` for LoRA fine-tuning
+- `tuner_backend`: LoRA backend, set to `peft`
+- `lora_rank`: LoRA rank, default `8`
+- `lora_alpha`: LoRA alpha coefficient, default `16`
+- `lora_dropout`: LoRA dropout ratio, default `0.05`
+
+**Dataset:**
+
+- `dataset`: Dataset path, supports local jsonl files (sharegpt format)
+- `max_length`: Maximum sequence length; sequences exceeding this will be truncated
+- `truncation_strategy`: Truncation strategy, options: `delete` (discard overlong samples) or `truncation_left`
+- `lazy_tokenize`: Whether to use lazy tokenization, `true` is recommended
+
+**Output:**
+
+- `output_dir`: Output directory
+- `save_steps`: Number of steps between saving checkpoints
+- `save_total_limit`: Maximum number of checkpoints to keep
+- `logging_steps`: Number of steps between logging
+- `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
+- `learning_rate`: Maximum learning rate; `1.0e-5` recommended for full fine-tuning, `3.0e-4` for LoRA fine-tuning
+- `num_train_epochs`: Number of training epochs
+- `lr_scheduler_type`: Learning rate scheduler type, `cosine` is recommended
+- `warmup_ratio`: Proportion of total training steps used for warmup
+- `bf16`: Whether to use BFloat16 mixed precision training
+
+**DeepSpeed / Optimization:**
+
+- `deepspeed`: DeepSpeed strategy, options: `zero0`, `zero2`, `zero2_offload`, `zero3`, `zero3_offload`; `zero3_offload` recommended for full fine-tuning, `zero2_offload` for LoRA fine-tuning
+- `gradient_checkpointing`: Whether to enable gradient checkpointing
+- `max_grad_norm`: Gradient clipping threshold
+
+**Other:**
+
+- `ddp_timeout`: Distributed training timeout (milliseconds)
+- `seed`: Random seed
+- `resume_from_checkpoint`: Resume training from a specified checkpoint path
+
+##### 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 in the `sft_train.sh` script:
+
+```shell
+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"}
+```
+
+Then, on each machine, execute the launch script in the `train/ms_swift_support/` directory:
+
+```sh
+# Single-machine training
+bash sft_train.sh
+
+# Multi-machine training (execute on each machine)
+IP_LIST="10.0.0.1,10.0.0.2" bash sft_train.sh
+```

finetune/README_CN.md

@@ -0,0 +1,416 @@
+<p align="left">
+   <a href="README.md">English</a> ｜ 中文
+</p>
+
+# 模型训练
+
+Hy3 提供了模型训练相关流程，您可以在此章节对训练数据格式进行处理以供模型训练使用。
+
+## 训练数据格式及处理
+
+**Hy3 同时支持慢思考与快思考两种模式，模型的默认输出是慢思考模式，若想让模型进行快思考，可通过 `reasoning_effort` 参数控制（可选值：`high`、`low`、`no_think`）。**
+
+训练数据按照以下形式处理为 messages 格式，训练和推理的默认 system prompt 为空，可以根据自己的需求进行设定。
+
+```python
+# Fast thinking pattern (no_think)
+{"reasoning_effort": "no_think", "messages": [{"content": "你是一个有用的人工智能助手。\n现在的时间是2026-01-01 13:26:12 周四", "role": "system"}, {"content": "1+1=?", "role": "user"}, {"role": "assistant", "content": "1+1=2"}]}
+
+# Slow thinking pattern (high)
+{"reasoning_effort": "high", "messages": [{"content": "你是一个有用的人工智能助手。\n现在的时间是2026-01-01 13:26:12 周四", "role": "system"}, {"content": "1+1=?", "role": "user"}, {"role": "assistant", "content": "1+1=2", "reasoning_content": "用户问的是1+1等于多少。在基本的十进制算术中，1+1等于2。"}]}
+
+from transformers import AutoTokenizer
+tokenizer = AutoTokenizer.from_pretrained("./models", use_fast=False, trust_remote_code=True)
+ids = tokenizer.apply_chat_template(messages, is_training=True)
+```
+
+## 权重格式转换
+
+Hy3 的原始 checkpoint 采用每个 expert 独立存储的格式，在训练前建议转换为 expert 融合后的 HuggingFace 标准格式（将同一层的多个 expert 权重融合为 3D 张量，并统一 key 命名），用于提高加载和训练的速率。不进行转换也可以直接使用原始格式进行训练，但加载速度会较慢。我们提供了转换脚本 `convert_ckpt_to_outer.py` 和校验脚本 `check_converted.py`，位于 `train/tools` 目录下。
+
+### 转换
+
+```sh
+python convert_ckpt_to_outer.py \
+    --input_dir <原始checkpoint目录> \
+    --output_dir <输出目录> \
+    --workers 8
+```
+
+**参数说明：**
+
+- `--input_dir`：原始 checkpoint 目录路径（必选）
+- `--output_dir`：转换后的 checkpoint 输出目录路径（必选）
+- `--workers`：并行转换的进程数，默认为 8（可选）
+
+转换脚本会执行以下步骤：
+1. 预扫描 `model.safetensors.index.json`，检测跨 shard 的 expert 分组
+2. 逐 shard 并行转换权重（key 重命名 + expert 融合）
+3. 后处理跨 shard 的 expert 分组（合并来自多个 shard 的数据）
+4. 复制 `config.json`、tokenizer 等其他文件
+5. 重建 `model.safetensors.index.json`
+
+### 校验
+
+转换完成后，建议使用校验脚本验证转换结果的完整性：
+
+```sh
+python check_converted.py <转换后的checkpoint目录> --spot-check 3
+```
+
+**参数说明：**
+
+- 第一个参数：转换后的 checkpoint 目录路径（必选）
+- `--spot-check`：随机抽检的 shard 文件数量，会加载 tensor 并检查 shape、dtype、NaN/Inf 等，默认为 3（可选）
+
+校验脚本会检查以下内容：
+1. `config.json` 的完整性
+2. `model.safetensors.index.json` 中所有预期 key 是否齐全（包括常规层和 MTP 层）
+3. 所有引用的 shard 文件是否存在且非空
+4. 抽检 shard 文件中 tensor 的 shape、dtype 是否正确，是否存在 NaN/Inf
+5. 检测孤立的空 shard 文件（跨 shard 合并残留，可安全删除）
+
+## 快速开始
+
+您可以参照快速开始文档中的内容进行快速上手。
+
+## 模型训练
+
+### 硬件需求
+
+经过测试，不开 make_moe_param_leaf_module 以及 zero3+offload，max_seq_length 为 4096：
+
+- **LoRA 微调**：最少需要单机 8 卡（显存至少 80GB）。
+- **全量微调**：最少需要 4 机 32 卡（显存至少 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` 目录下
+- **ms-swift 训练**：位于 `train/ms_swift_support` 目录下
+
+#### DeepSpeed 原生训练
+
+参考：[HuggingFace Transformers Trainer](https://huggingface.co/docs/transformers/main/en/main_classes/trainer)
+
+##### 单机启动训练
+
+在 `train/deepspeed_support` 目录下，执行：
+
+```sh
+pip install -r requirements.txt
+bash train.sh
+```
+
+##### 多机启动训练
+
+如果要用多台机器启动训练，请先完成 [配置机器间免密 ssh 登录](#配置机器间免密-ssh-登录多机训练) 中的配置，并保证多台机器在一个集群内。
+
+确认依赖已经安装完成（如未安装，请执行`pip install -r requirements.txt`安装），然后在`train.sh`中的开头增加以下配置：
+
+```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/`目录下，执行`bash train.sh`即可，注意第一次启动时可能会看见以下的输出：
+
+```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`: 要加载的 Hy3 的 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 ，根据网站的指引进行安装。
+
+##### 配置文件
+
+我们提供了 llama-factory 的训练示例配置文件 `hy_v3_lora_sft.yaml`和`hy_v3_full_sft.yaml`文件，分别对应 LoRA 训练和全量微调。
+
+脚本中的关键参数如下：
+
+**模型相关：**
+
+- `model_name_or_path`: Hy3 HF 格式预训练模型权重路径
+- `trust_remote_code`: 是否信任远程代码, Hy3 需要设置为 `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`: 对话模板, Hy3 使用 `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-登录多机训练) 中的配置（单机训练可跳过此步骤）。
+
+修改`train_lf.sh`中开头的以下配置：
+
+```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/`目录下执行`bash train_lf.sh`。
+
+#### ms-swift 训练
+
+如果对 ms-swift 较为熟悉，可使用 ms-swift 进行微调。脚本、代码以及配置文件都归档在 `train/ms_swift_support` 目录下。如果没有特别说明，接下来我们提到的文件都是该目录下的文件。
+
+##### 安装
+
+可以通过 pip 安装 ms-swift：
+
+```sh
+pip install ms-swift==4.2.2
+```
+
+或从源码安装：https://github.com/modelscope/ms-swift
+
+##### 训练脚本与配置文件
+
+| 训练方式 | 配置文件 | 启动脚本 |
+|---------|---------|---------|
+| 全量微调 | `hy_v3_full_sft.yaml` | `bash sft_train.sh` |
+| LoRA 微调 | `hy_v3_lora_sft.yaml` | `bash sft_train.sh` |
+
+##### 关于 eos_token_id Patch
+
+目录下的 `hy_v3_swift_patches.py` 文件用于修复 ms-swift 默认模板中 eos token 的问题。默认模板将 `<｜hy_eos｜>` 字符串作为 `chat_sep` 和 `suffix`，该字符串会被 tokenize 为多个 token ID，导致推理时 `model.generate()` 无法正确停止。
+
+Patch 通过 `[['eos_token_id']]` 语法重新注册模板，使 ms-swift 在运行时动态解析 `tokenizer.eos_token_id`，生成正确的单个 token。
+
+启动脚本已通过 `--custom_register_path hy_v3_swift_patches.py` 自动加载此 patch，无需额外操作。
+
+##### 关键参数
+
+配置文件中的关键参数如下：
+
+**模型相关：**
+
+- `model`: 模型路径，可以是 HuggingFace Hub ID 或本地路径
+- `model_type`: 模型类型，设为 `hy_v3`
+- `template`: 对话模板，设为 `hy_v3`
+- `torch_dtype`: 数据类型，推荐 `bfloat16`
+- `attn_impl`: 注意力实现，推荐 `flash_attn`
+
+**训练方法：**
+
+- `tuner_type`: 微调类型，全量微调设为 `full`，LoRA 微调设为 `lora`
+- `tuner_backend`: LoRA 后端，设为 `peft`
+- `lora_rank`: LoRA 秩，默认 `8`
+- `lora_alpha`: LoRA alpha 系数，默认 `16`
+- `lora_dropout`: LoRA dropout 比率，默认 `0.05`
+
+**数据集：**
+
+- `dataset`: 数据集路径，支持本地 jsonl 文件（sharegpt 格式）
+- `max_length`: 最大序列长度，超出部分会被截断
+- `truncation_strategy`: 截断策略，可选 `delete`（丢弃超长样本）或 `truncation_left`
+- `lazy_tokenize`: 是否延迟 tokenize，推荐 `true`
+
+**输出：**
+
+- `output_dir`: 输出目录
+- `save_steps`: 每多少步保存一次 checkpoint
+- `save_total_limit`: 最多保留的 checkpoint 数量
+- `logging_steps`: 每多少步记录一次日志
+- `report_to`: 日志上报工具，可选 `none`, `wandb`, `tensorboard`, `swanlab`, `mlflow`
+
+**训练超参数：**
+
+- `per_device_train_batch_size`: 每张卡上的 batch size
+- `gradient_accumulation_steps`: 梯度累积步数
+- `learning_rate`: 最大学习率，全量微调推荐 `1.0e-5`，LoRA 微调推荐 `3.0e-4`
+- `num_train_epochs`: 训练轮数
+- `lr_scheduler_type`: 学习率调度器类型，推荐 `cosine`
+- `warmup_ratio`: 预热阶段占总训练步数的比例
+- `bf16`: 是否使用 BFloat16 混合精度训练
+
+**DeepSpeed / 优化：**
+
+- `deepspeed`: DeepSpeed 策略，可选 `zero0`, `zero2`, `zero2_offload`, `zero3`, `zero3_offload`；全量微调推荐 `zero3_offload`，LoRA 微调推荐 `zero2_offload`
+- `gradient_checkpointing`: 是否开启梯度重计算
+- `max_grad_norm`: 梯度裁剪阈值
+
+**其他：**
+
+- `ddp_timeout`: 分布式训练超时时间（毫秒）
+- `seed`: 随机种子
+- `resume_from_checkpoint`: 从指定 checkpoint 路径恢复训练
+
+##### 启动训练
+
+如需多机训练，请先完成 [配置机器间免密 ssh 登录](#配置机器间免密-ssh-登录多机训练) 中的配置（单机训练可跳过此步骤）。
+
+修改 `sft_train.sh` 脚本中的以下配置：
+
+```shell
+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"}
+```
+
+然后，在每一台机器上，在 `train/ms_swift_support/` 目录下执行启动脚本：
+
+```sh
+# 单机训练
+bash sft_train.sh
+
+# 多机训练（在每台机器上执行）
+IP_LIST="10.0.0.1,10.0.0.2" bash sft_train.sh
+```

finetune/data/example_data.jsonl

@@ -0,0 +1,8 @@
+{"messages": [{"role": "system", "content": "You are a helpful assistant."}, {"role": "user", "content": "能否帮我撰写一个关于环保的议论文开头部分？\n主题：环保\n立场：重要性和我们的责任"}, {"role": "assistant", "content": "在我们的生活中，环保已经成为一个不可忽视的议题。随着科技的发展和人口的增长，我们的环境正在遭受前所未有的压力。森林被砍伐，河流被污染，动植物种类正在消失，全球气候正在发生变化。这些都是我们无法忽视的事实。然而，环保并不仅仅是一个环境问题，它更是一个关乎我们生存和发展的问题。环保的重要性已经超越了任何一个单一的领域，它关乎我们的健康、经济、社会和文化。因此，我们每一个人都有责任和义务去保护我们的环境，这不仅是为了我们自己，也是为了我们的后代和整个地球。"}]}
+{"messages": [{"role": "system", "content": "You are a helpful assistant."}, {"role": "user", "content": "请描述小红书APP的用户画像，并用表格的形式展示结果。"}, {"role": "assistant", "content": "| 用户特征   | 描述                           |\n| -------- | ------------------------------ |\n| 年龄段    | 18-35岁                        |\n| 性别      | 女性用户较多（约占70%）          |\n| 地域      | 一线城市和二线城市用户占比较高        |\n| 职业      | 白领、创业者、自由职业者等           |\n| 爱好      | 美容、时尚、旅行、美食、健康、生活方式等 |\n| 消费能力  | 中高消费能力用户                  |\n| 使用目的  | 阅读内容，分享购物心得，结识相似兴趣的朋友  |\n\n注：以上用户画像仅供参考，实际情况可能会有所不同。"}]}
+{"messages": [{"role": "system", "content": "You are a helpful assistant."}, {"role": "user", "content": "请根据以下两个句子的描述，判断它们是否对同一情况进行了描述。句子1：“外面下着大雨，人们纷纷打起了伞。” 句子2：“由于雨量较大，人们都撑起了雨伞。”"}, {"role": "assistant", "content": "是的，这两个句子描述的是同一情况。"}]}
+{"messages": [{"role": "system", "content": "You are a helpful assistant."}, {"role": "user", "content": "你认为游戏开发者应该如何处理游戏中的社交网络集成？"}, {"role": "assistant", "content": "游戏开发者在处理游戏中的社交网络集成时可以考虑以下几点：\n\n1. 以用户为中心：确保集成社交网络功能是为了提高用户体验，让玩家更容易地与朋友互动、分享成就和游戏进度。\n\n2. 保护隐私：在集成社交网络时，务必保护用户的隐私。在收集和使用用户数据时，需要遵循相关法规和用户协议，同时提供明确的隐私政策。\n\n3. 适度推送：在推送游戏动态、邀请和通知时要适度，避免过度打扰用户或侵犯用户隐私。可提供自定义设置，让玩家选择性地接收通知。\n\n4. 社交功能设计：在游戏中加入社交功能，如聊天、组队、比赛等，以增进玩家之间的互动。同时，确保这些功能易用且不影响游戏体验。\n\n5. 社区建设：鼓励玩家在社交平台上建立游戏相关的社区，分享游戏心得、攻略和经验。这将有助于增加游戏的知名度和吸引力。\n\n6. 奖励机制：设定奖励机制，鼓励玩家在社交网络上分享游戏内容、邀请好友等。但要确保奖励机制不会破坏游戏平衡。\n\n7. 持续更新和优化：根据玩家的反馈，不断优化社交网络集成功能，确保其与游戏的其他方面保持一致。\n\n总之，游戏开发者应该在保护用户隐私的前提下，提供便捷的社交网络集成功能，以提高游戏的互动性和吸引力。同时，要关注玩家的反馈，不断优化功能。"}]}
+{"messages": [{"role": "system", "content": "You are a helpful assistant."}, {"role": "user", "content": "将这段过去进行时的文字改为现在完成进行时：“她们在公园里跳舞。”"}, {"role": "assistant", "content": "她们已经在公园里跳舞了。"}]}
+{"messages": [{"role": "system", "content": "You are a helpful assistant."}, {"role": "user", "content": "在godaddy上买了一台服务器怎么更改root账户"}, {"role": "assistant", "content": "在GoDaddy上购买的服务器可能是Linux系统的VPS或者是专用服务器。要更改root账户，你可以通过SSH连接到服务器并修改密码。以下是具体步骤：\n\n1. 打开终端（Terminal）或使用SSH客户端（如PuTTY）。\n2. 使用SSH连接到服务器。在终端中输入以下命令：\n\n   ```\n   ssh root@your_server_ip_address\n   ```\n\n   将`your_server_ip_address`替换为你的服务器IP地址。如果需要，你可能还需要输入服务器的SSH端口号。\n   \n3. 当系统提示输入密码时，输入你的当前root密码。如果连接成功，你将看到服务器的命令提示符。\n\n4. 要更改root密码，请输入以下命令：\n\n   ```\n   passwd\n   ```\n\n5. 系统将提示你输入新的UNIX密码。输入你想要设置的新密码并按回车键。\n\n6. 系统将要求你重新输入新密码以进行确认。再次输入新密码并按回车键。\n\n7. 如果密码更改成功，系统将显示“密码已成功更新”的消息。\n\n8. 使用`exit`命令断开与服务器的连接。\n\n现在你已经成功更改了root账户的密码。下次使用SSH连接到服务器时，请使用新的密码进行身份验证。"}]}
+{"messages": [{"role": "system", "content": "You are a helpful assistant."}, {"role": "user", "content": "PHP算法逻辑：如何计算年龄？"}, {"role": "assistant", "content": "要计算年龄，您可以使用以下PHP算法。此算法考虑了生日、当前日期以及闰年的计算。\n\n```php\nfunction calculateAge($birthdate) {\n  // 格式为YYYY-MM-DD的生日\n  $birthdate = strtotime($birthdate);\n  \n  // 获取当前日期\n  $now = time();\n\n  // 计算年龄\n  $age = date('Y', $now) - date('Y', $birthdate);\n  \n  // 如果还没有过生日，减去一年\n  if (date('m-d', $now) < date('m-d', $birthdate)) {\n    $age--;\n  }\n\n  return $age;\n}\n\n// 使用\n$birthdate = \"2000-01-01\";\necho \"年龄是: \" . calculateAge($birthdate) . \" 岁\";\n```\n\n这个函数将生日作为输入，输出用户的年龄。首先，它将生日转换为时间戳，然后获取当前的时间戳。接下来，它计算年份之间的差异。如果当前日期仍未超过这一年的生日，这个函数将年龄减少1。\n\n在使用函数时，只需传入出生日期即可，函数会返回相应的年龄。"}]}
+{"messages": [{"role": "system", "content": "You are a helpful assistant."}, {"role": "user", "content": "给我解释一下什么是胰岛素抵抗？ 请详细描述胰岛素抵抗的症状、原因与治疗方法。 胰岛素抵抗会影响一个人的日常生活吗？ 有哪些生活习惯改变可以帮助改善胰岛素抵抗？"}, {"role": "assistant", "content": "胰岛素抵抗是指身体对胰岛素的敏感性降低，胰岛素不能有效地发挥作用，导致血糖无法正常进入身体细胞。胰岛素是一种激素，由胰腺产生，其主要作用是帮助身体细胞吸收血糖并将其转化为能量。当胰岛素抵抗发生时，胰腺会分泌更多的胰岛素以维持正常的血糖水平，但长期的胰岛素抵抗可能导致2型糖尿病和其他健康问题的发生。\n\n胰岛素抵抗的症状可能包括：\n\n1. 高血糖：由于胰岛素不能正常工作，血糖无法进入细胞，导致血糖升高。\n2. 疲劳：由于身体无法有效利用血糖产生能量，患者可能会感到疲劳无力。\n3. 肥胖：胰岛素抵抗可能导致过量的胰岛素分泌，使得脂肪存储增加，进而导致体重增加。\n4. 高血压：胰岛素抵抗与高血压之间存在关联，可能导致血压升高。\n5. 黑棘皮病：皮肤出现黑色、粗糙的皮疹，通常出现在颈部、腋窝等皮肤皱褶处。\n\n胰岛素抵抗的原因可能包括：\n\n1. 遗传因素：遗传易感基因可能会增加一个人患胰岛素抵抗的风险。\n2. 肥胖：过重或肥胖会增加患胰岛素抵抗的风险。\n3. 缺乏运动：缺乏运动可能会导致肌肉对胰岛素的敏感性降低。\n4. 高糖饮食：长期摄入高糖食品可能导致胰岛素抵抗。\n5. 睡眠不足：长期睡眠不足可能会影响胰岛素的正常功能。\n\n治疗胰岛素抵抗的方法主要包括生活方式的改变和药物治疗：\n\n1. 减轻体重：通过改善饮食和增加运动来减轻体重，有助于提高身体对胰岛素的敏感性。\n2. 增加运动：增加有氧运动和力量训练，有助于提高肌肉对胰岛素的敏感性。\n3. 改善饮食：摄入低糖、高纤维、富含全谷物和蔬菜水果的食物，有助于控制血糖和减轻体重。\n4. 药物治疗：在必要时，医生可能会开处降糖药物，如二甲双胍，以帮助提高胰岛素敏感性。\n\n胰岛素抵抗会影响一个人的日常生活，包括疲劳、体重增加、血糖波动等问题。通过改变生活习惯，如减轻体重、增加运动、改善饮食和保证充足睡眠，可以帮助改善胰岛素抵抗，降低患2型糖尿病的风险。"}]}

finetune/deepspeed_support/ds_zero2_no_offload.json

@@ -0,0 +1,25 @@
+{
+    "fp16": {
+        "enabled": "auto",
+        "loss_scale": 0,
+        "loss_scale_window": 100,
+        "initial_scale_power": 16,
+        "hysteresis": 2,
+        "min_loss_scale": 1e-10
+    },
+    "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
+}

finetune/deepspeed_support/ds_zero2_offload.json

@@ -0,0 +1,34 @@
+{
+    "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": 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_batch_size": "auto",
+    "train_micro_batch_size_per_gpu": "auto",
+    "wall_clock_breakdown": false
+}

finetune/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
+}

finetune/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
+}

finetune/deepspeed_support/ds_zero3_optimizer_offload.json

@@ -0,0 +1,37 @@
+{
+    "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
+        },
+        "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
+}

finetune/deepspeed_support/ds_zero3_param_offload.json

@@ -0,0 +1,37 @@
+{
+    "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_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": 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
+}

finetune/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()

finetune/deepspeed_support/merge_lora_weight.sh

@@ -0,0 +1 @@
+python3 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

finetune/deepspeed_support/train.py

@@ -0,0 +1,903 @@
+# 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'
+        try:
+            template_output = self.tokenizer.apply_chat_template(data_dict['messages'], tokenize=True, return_dict=False, is_training=True, reasoning_effort=reasoning_effort)
+        except Exception as e:
+            print(f"[ERROR] apply_chat_template failed: {e}")
+            print(f"[ERROR] messages: {data_dict['messages']}")
+            print(f"[ERROR] reasoning_effort: {reasoning_effort}")
+            template_output = []
+        
+        # Debug: Check template_output type and content
+        if isinstance(template_output, bool):
+            print(f"[WARNING] apply_chat_template returned bool: {template_output}")
+            print(f"[WARNING] messages: {data_dict['messages']}")
+            print(f"[WARNING] reasoning_effort: {reasoning_effort}")
+            # Return empty tensor to avoid crash
+            template_output = []
+        
+        if isinstance(template_output, list) and len(template_output) > 0 and isinstance(template_output[0], list):
+            template_output = template_output[0]
+        
+        # Ensure template_output is a list of integers
+        if not isinstance(template_output, list) or not all(isinstance(x, int) for x in template_output):
+            print(f"[WARNING] Invalid template_output format: {type(template_output)}, content: {template_output}")
+            print(f"[WARNING] messages: {data_dict['messages']}")
+            template_output = []
+        
+        message_tokens = torch.tensor(template_output, dtype=torch.long)
+
+        # Use new HunYuan tokenizer special tokens
+        # Get assistant_token from tokenizer attribute (dynamic, not hardcoded)
+        assistant_token = getattr(self.tokenizer, 'assistant_token', None)
+        if assistant_token is None:
+            # Fallback: try to get from tokenizer_config
+            assistant_token = '<｜hy_Assistant:6124c78e｜>'
+        assistant_token_id = self.tokenizer.convert_tokens_to_ids(assistant_token)
+        
+        # Safety check: ensure assistant_token_id is valid
+        if assistant_token_id is None or assistant_token_id == self.tokenizer.unk_token_id:
+            print(f"[WARNING] assistant_token_id is invalid: {assistant_token_id}, assistant_token: {assistant_token}")
+            print(f"[WARNING] Using fallback token ID")
+            # Use a fallback: try to find the token in vocab
+            assistant_token_id = self.tokenizer.convert_tokens_to_ids('<｜hy_Assistant:6124c78e｜>')
+        
+        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 assistant_token, ends at eos_token
+        # Handle empty message_tokens case
+        if message_tokens.numel() == 0:
+            print(f"[WARNING] Empty message_tokens, skipping data sample")
+            # Return empty tensors with proper shape
+            input_ids = torch.tensor([], dtype=torch.long)
+            labels = torch.tensor([], dtype=torch.long)
+            attention_mask = torch.tensor([], dtype=torch.bool)
+        else:
+            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]
+            attention_mask = torch.tensor(attention_mask, dtype=torch.bool)
+
+        model_inputs["input_ids"] = input_ids
+        model_inputs["attention_mask"] = attention_mask
+        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)
+        
+        # Check if the encoded data is empty (due to tokenization failure)
+        if model_inputs["input_ids"].numel() == 0:
+            # Return a valid placeholder sample to avoid crash
+            # Use a minimal valid sequence with 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
+            
+            # Create a minimal valid sequence: <｜hy_Assistant｜> + eos
+            placeholder_tokens = [assistant_token_id, eos_token_id]
+            placeholder_tokens = placeholder_tokens[:self.max_seq_length]
+            
+            input_ids = torch.tensor(placeholder_tokens, dtype=torch.long)
+            labels = torch.tensor([IGNORE_INDEX, eos_token_id], dtype=torch.long)[:self.max_seq_length]
+            attention_mask = torch.tensor([1, 1], dtype=torch.bool)[:self.max_seq_length]
+            
+            # Pad to max_seq_length if needed
+            if len(placeholder_tokens) < self.max_seq_length:
+                padding_length = self.max_seq_length - len(placeholder_tokens)
+                input_ids = torch.cat([input_ids, torch.full((padding_length,), pad_token_id, dtype=torch.long)])
+                labels = torch.cat([labels, torch.full((padding_length,), IGNORE_INDEX, dtype=torch.long)])
+                attention_mask = torch.cat([attention_mask, torch.zeros(padding_length, dtype=torch.bool)])
+            
+            model_inputs = {
+                "input_ids": input_ids,
+                "attention_mask": attention_mask,
+                "labels": labels
+            }
+
+        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.
+    # -----------------------------------------------------------------------
+    # Key renames: checkpoint format -> installed transformers 5.8.1 model format
+    # Checkpoint uses: mlp.router.gate, mlp.expert_bias, mlp.shared_mlp
+    # Model uses:      mlp.gate,        mlp.e_score_correction_bias, mlp.shared_experts
+    _CKPT_KEY_RENAMES = [
+        ("mlp.router.gate.", "mlp.gate."),
+        ("mlp.expert_bias", "mlp.e_score_correction_bias"),
+        ("mlp.shared_mlp.", "mlp.shared_experts."),
+        # Also handle even older checkpoints that use mlp.gate.wg
+        ("mlp.gate.wg.", "mlp.gate."),
+    ]
+
+    # Regex to match per-expert keys in checkpoint
+    # e.g. model.layers.10.mlp.experts.5.gate_proj.weight
+    _EXPERT_KEY_RE = re.compile(
+        r"^(.*\.mlp\.experts\.)(\d+)\.(gate_proj|up_proj|down_proj)\.weight$"
+    )
+
+    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 = {}
+        expert_groups = {}  # prefix -> {expert_idx -> {proj_name -> tensor}}
+
+        for k, v in state_dict.items():
+            m = _EXPERT_KEY_RE.match(k)
+            if m:
+                # Per-expert key: collect for fusion
+                prefix = m.group(1)
+                expert_idx = int(m.group(2))
+                proj_name = m.group(3)
+                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] = v
+            else:
+                # Non-expert key: apply simple renames
+                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
+
+        # Fuse expert groups into 3D tensors
+        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:
+                    continue
+                exp = experts[i]
+                if "gate_proj" in exp and "up_proj" in exp:
+                    gate_up_list.append(torch.cat([exp["gate_proj"], exp["up_proj"]], dim=0))
+                if "down_proj" in exp:
+                    down_list.append(exp["down_proj"])
+            if gate_up_list:
+                new_sd[f"{prefix}gate_up_proj"] = torch.stack(gate_up_list, dim=0)
+            if down_list:
+                new_sd[f"{prefix}down_proj"] = torch.stack(down_list, dim=0)
+        del expert_groups
+
+        # 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}")
+        # ---------------------------------------------------------------
+        # Memory-efficient loading: Instead of from_pretrained's default
+        # ZeRO-3 path (which merges ALL shards into one huge dict in CPU
+        # memory), we:
+        #   1. Create the model skeleton under deepspeed.zero.Init (meta)
+        #   2. Load each safetensors shard one at a time
+        #   3. Scatter each shard's weights into ZeRO-3 partitions
+        #   4. Free the shard immediately
+        # This reduces per-rank CPU memory from ~670GB to ~7GB (1 shard).
+        # ---------------------------------------------------------------
+        import json as _json
+        from safetensors import safe_open
+
+        ds_config = training_args.deepspeed
+        if isinstance(ds_config, str):
+            with open(ds_config, "r") as f:
+                ds_config = _json.load(f)
+        # Replace "auto" values that deepspeed.zero.Init cannot resolve
+        _auto_defaults = {
+            "train_batch_size": training_args.per_device_train_batch_size
+                               * training_args.gradient_accumulation_steps
+                               * training_args.world_size,
+            "train_micro_batch_size_per_gpu": training_args.per_device_train_batch_size,
+            "gradient_accumulation_steps": training_args.gradient_accumulation_steps,
+            "gradient_clipping": training_args.max_grad_norm,
+        }
+        for k, v in _auto_defaults.items():
+            if k in ds_config and ds_config[k] == "auto":
+                ds_config[k] = v
+
+        # Step 1: Create model with empty (meta) weights under ZeRO-3 Init
+        model_path = training_args.model_name_or_path
+        config = transformers.AutoConfig.from_pretrained(
+            model_path, trust_remote_code=True
+        )
+        with deepspeed.zero.Init(dtype=torch.bfloat16,
+                                 config_dict_or_path=ds_config):
+            model = transformers.AutoModelForCausalLM.from_config(
+                config, trust_remote_code=True,
+                torch_dtype=init_kwargs.get("dtype", torch.bfloat16),
+                attn_implementation=init_kwargs.get("attn_implementation", None),
+            )
+        print(f"[HYV3] Model skeleton created under ZeRO-3 Init.")
+
+        # Step 2: Determine shard files from index
+        index_file = os.path.join(model_path, "model.safetensors.index.json")
+        if os.path.isfile(index_file):
+            with open(index_file, "r") as f:
+                index_data = _json.load(f)
+            # Get unique shard filenames in order
+            shard_files = list(dict.fromkeys(index_data["weight_map"].values()))
+        else:
+            # Single shard model
+            shard_files = ["model.safetensors"]
+
+        # Step 3: Load each shard and scatter into ZeRO-3 model
+        # For per-expert keys, we need to collect them per-layer and fuse
+        # into 3D tensors (gate_up_proj, down_proj) before scattering.
+        total_shards = len(shard_files)
+        all_loaded_keys = set()
+        # Buffer for cross-shard expert accumulation:
+        # prefix -> {expert_idx -> {proj_name -> tensor}}
+        pending_experts = {}
+
+        for shard_idx, shard_name in enumerate(shard_files, 1):
+            shard_path = os.path.join(model_path, shard_name)
+            print(f"[HYV3] Loading shard {shard_idx}/{total_shards}: {shard_name}")
+
+            # Load shard into CPU memory
+            shard_sd = {}
+            with safe_open(shard_path, framework="pt", device="cpu") as f:
+                for key in f.keys():
+                    shard_sd[key] = f.get_tensor(key)
+
+            # Separate expert keys from non-expert keys, apply renames
+            renamed_sd = {}
+            expert_keys_in_shard = {}  # prefix -> {expert_idx -> {proj_name -> tensor}}
+
+            for k, v in shard_sd.items():
+                m = _EXPERT_KEY_RE.match(k)
+                if m:
+                    # Per-expert key: collect for fusion
+                    prefix = m.group(1)  # e.g. "model.layers.10.mlp.experts."
+                    expert_idx = int(m.group(2))
+                    proj_name = m.group(3)  # gate_proj, up_proj, or down_proj
+                    if prefix not in expert_keys_in_shard:
+                        expert_keys_in_shard[prefix] = {}
+                    if expert_idx not in expert_keys_in_shard[prefix]:
+                        expert_keys_in_shard[prefix][expert_idx] = {}
+                    expert_keys_in_shard[prefix][expert_idx][proj_name] = v
+                else:
+                    # Non-expert key: apply simple renames
+                    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
+                    renamed_sd[new_k] = v
+            del shard_sd
+
+            # Merge expert keys from this shard into pending_experts
+            for prefix, experts in expert_keys_in_shard.items():
+                if prefix not in pending_experts:
+                    pending_experts[prefix] = {}
+                for idx, projs in experts.items():
+                    if idx not in pending_experts[prefix]:
+                        pending_experts[prefix][idx] = {}
+                    pending_experts[prefix][idx].update(projs)
+            del expert_keys_in_shard
+
+            # Check if any pending expert groups are now complete
+            # (all 3 projections for all experts in the layer)
+            # We detect completeness by checking if we have gate_proj, up_proj,
+            # and down_proj for a contiguous range of expert indices.
+            completed_prefixes = []
+            for prefix, experts in pending_experts.items():
+                # Check if all experts have all 3 projections
+                if not experts:
+                    continue
+                max_idx = max(experts.keys())
+                num_experts_found = len(experts)
+                # A layer is complete if we have a contiguous range and all have 3 projs
+                all_complete = all(
+                    len(projs) == 3 for projs in experts.values()
+                )
+                # Heuristic: if we have 192 experts (or max_idx+1 == num found)
+                # and all have 3 projections, consider it complete
+                if all_complete and num_experts_found == (max_idx + 1):
+                    completed_prefixes.append(prefix)
+
+            # Fuse completed expert groups and add to renamed_sd
+            for prefix in completed_prefixes:
+                experts = pending_experts.pop(prefix)
+                num_experts_layer = max(experts.keys()) + 1
+                gate_up_list = []
+                down_list = []
+                for i in range(num_experts_layer):
+                    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)
+                del gate_up_list, down_list, experts
+
+                # Model key format: model.layers.X.mlp.experts.gate_up_proj
+                renamed_sd[f"{prefix}gate_up_proj"] = fused_gate_up
+                renamed_sd[f"{prefix}down_proj"] = fused_down
+                print(f"[HYV3]   Fused {num_experts_layer} experts for {prefix}")
+
+            # Scatter this shard's weights into ZeRO-3 partitioned model
+            if renamed_sd:
+                _orig_load_zero3(model, renamed_sd)
+
+                # Also load buffers (e.g. e_score_correction_bias)
+                for name, buf in model.named_buffers():
+                    if name in renamed_sd:
+                        src_tensor = renamed_sd[name]
+                        if isinstance(src_tensor, torch.Tensor):
+                            buf.data.copy_(src_tensor.to(buf.dtype))
+
+                all_loaded_keys.update(renamed_sd.keys())
+            del renamed_sd
+            import gc; gc.collect()
+
+        # Flush any remaining pending experts (cross-shard edge case)
+        if pending_experts:
+            print(f"[HYV3] Flushing {len(pending_experts)} remaining expert group(s)...")
+            flush_sd = {}
+            for prefix, experts in pending_experts.items():
+                num_experts_layer = max(experts.keys()) + 1
+                gate_up_list = []
+                down_list = []
+                for i in range(num_experts_layer):
+                    if i not in experts:
+                        print(f"[HYV3] Warning: Missing expert {i} in {prefix}")
+                        continue
+                    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"])
+                if gate_up_list:
+                    fused_gate_up = torch.stack(gate_up_list, dim=0)
+                    fused_down = torch.stack(down_list, dim=0)
+                    flush_sd[f"{prefix}gate_up_proj"] = fused_gate_up
+                    flush_sd[f"{prefix}down_proj"] = fused_down
+                    print(f"[HYV3]   Fused {len(gate_up_list)} experts for {prefix}")
+                del gate_up_list, down_list
+            del pending_experts
+
+            if flush_sd:
+                _orig_load_zero3(model, flush_sd)
+                for name, buf in model.named_buffers():
+                    if name in flush_sd:
+                        src_tensor = flush_sd[name]
+                        if isinstance(src_tensor, torch.Tensor):
+                            buf.data.copy_(src_tensor.to(buf.dtype))
+                all_loaded_keys.update(flush_sd.keys())
+            del flush_sd
+            import gc; gc.collect()
+
+        # Step 4: Report any missing/unexpected keys
+        model_keys = set(n for n, _ in model.named_parameters())
+        model_keys.update(n for n, _ in model.named_buffers())
+        missing = model_keys - all_loaded_keys
+        unexpected = all_loaded_keys - model_keys
+        if missing:
+            # Filter out keys that are expected to be missing (e.g. lm_head with tied embeddings)
+            real_missing = {k for k in missing if "lm_head" not in k}
+            if real_missing:
+                print(f"[HYV3] Warning: {len(real_missing)} keys not found in checkpoint "
+                      f"(first 10): {list(real_missing)[:10]}")
+        if unexpected:
+            print(f"[HYV3] Warning: {len(unexpected)} unexpected keys in checkpoint "
+                  f"(first 10): {list(unexpected)[:10]}")
+        print(f"[HYV3] Shard-by-shard loading complete. "
+              f"Loaded {len(all_loaded_keys)} keys from {total_shards} shards.")
+    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)
+
+    # Synchronize all processes before exit to avoid "Connection reset by peer"
+    # warnings caused by timing differences in multi-node shutdown.
+    if torch.distributed.is_initialized():
+        torch.distributed.barrier()
+
+
+if __name__ == "__main__":
+    train()

finetune/deepspeed_support/train.sh

@@ -0,0 +1,118 @@
+#!/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}))
+
+model_path=path_to_model_weight
+tokenizer_path=../../models
+train_data_file=../data/example_data.jsonl
+
+# ds_config_file=ds_zero2_no_offload.json
+# ds_config_file=ds_zero3_no_offload.json
+ds_config_file=ds_zero3_offload.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} \
+    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 \
+    --hidden_size 4096 \
+    --intermediate_size 13312 \
+    --model_max_length 4096 \
+    --max_seq_length 4096 \
+    --moe_topk 8 \
+    --num_experts 192 \
+    --moe_intermediate_size 1536 \
+    --moe_layer_num_skipped 1 \
+    --num_attention_heads 64 \
+    --num_key_value_heads 8 \
+    --num_layers 80 \
+    --use_mixed_mlp_moe \
+    --num_shared_expert 1 \
+    --use_qk_norm | tee ${log_file}

finetune/deepspeed_support/train_lora.sh

@@ -0,0 +1,123 @@
+#!/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}))
+
+
+model_path=path_to_model_weight
+tokenizer_path=../../models
+train_data_file=../data/example_data.jsonl
+
+# ds_config_file=ds_zero2_no_offload.json
+# ds_config_file=ds_zero3_no_offload.json
+ds_config_file=ds_zero3_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} \
+    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 4096 \
+    --intermediate_size 13312 \
+    --model_max_length 4096 \
+    --max_seq_length 4096 \
+    --moe_topk 8 \
+    --num_experts 192 \
+    --moe_intermediate_size 1536 \
+    --moe_layer_num_skipped 1 \
+    --num_attention_heads 64 \
+    --num_key_value_heads 8 \
+    --num_layers 80 \
+    --use_mixed_mlp_moe \
+    --num_shared_expert 1 \
+    --use_qk_norm | tee ${log_file}

finetune/llama_factory_support/dataset_info.json

@@ -0,0 +1,16 @@
+{
+  "hy_v3_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"
+    }
+  }
+}

finetune/llama_factory_support/hy_v3_full_sft.yaml

@@ -0,0 +1,57 @@
+### model
+model_name_or_path: ../hf
+trust_remote_code: true
+
+### method
+stage: sft
+do_train: true
+finetuning_type: full
+deepspeed: ../deepspeed_support/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: 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: 1.0e-5
+num_train_epochs: 100.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
+
+### HYV3 specific parameters (MoE architecture)
+# moe_num_experts: 192
+# moe_topk: 8
+# moe_intermediate_size: 1536
+# moe_layer_num_skipped: 1
+# output_router_logits: true

finetune/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: ../deepspeed_support/ds_zero2_offload.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_demo
+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: 4
+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

finetune/llama_factory_support/hy_v3_patches.py

@@ -0,0 +1,621 @@
+"""
+HYV3 monkey-patches for LLaMA Factory + DeepSpeed ZeRO-3 training.
+
+This module applies all necessary runtime patches so that HYV3 (MoE)
+can be trained correctly under LLaMA Factory with DeepSpeed ZeRO-3.
+
+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. ZeRO-3 key rename + expert fuse + buffer loading
+       Handles both inner format (original) and outer format (pre-converted) checkpoints.
+       - Key renaming: mlp.router.gate -> mlp.gate, etc.
+       - Per-expert -> 3D fuse: experts.N.gate_proj -> experts.gate_up_proj
+       - Buffer loading: e_score_correction_bias (ZeRO-3 only handles parameters)
+    2. Tokenizer file copy (CustomSaveCallback)
+       Ensures each checkpoint directory is self-contained for inference.
+"""
+
+import os
+import re
+import logging
+import shutil
+from typing import Optional
+
+import torch
+
+logger = logging.getLogger(__name__)
+
+# ============================================================================
+# Patch 1: Key rename + expert fuse + buffer loading for ZeRO-3
+#
+# The checkpoint may be in either inner format (original) or outer format
+# (pre-converted by convert_ckpt_to_outer.py). This patch handles both:
+#   - Key renaming: mlp.router.gate -> mlp.gate, etc.
+#   - Per-expert -> 3D fuse: experts.N.gate_proj -> experts.gate_up_proj
+#   - Buffer loading: e_score_correction_bias (ZeRO-3 only handles parameters)
+#
+# If the checkpoint is already in outer format, the rename/fuse logic is
+# effectively a no-op (no matching keys to transform).
+# ============================================================================
+
+# Key renames: checkpoint inner format -> model format
+_CKPT_KEY_RENAMES = [
+    ("mlp.router.gate.", "mlp.gate."),
+    ("mlp.expert_bias", "mlp.e_score_correction_bias"),
+    ("mlp.shared_mlp.", "mlp.shared_experts."),
+    # Also handle even older checkpoints that use mlp.gate.wg
+    ("mlp.gate.wg.", "mlp.gate."),
+]
+
+# Regex to match per-expert keys in checkpoint
+# e.g. model.layers.10.mlp.experts.5.gate_proj.weight
+_EXPERT_KEY_RE = re.compile(
+    r"^(.*\.mlp\.experts\.)(\d+)\.(gate_proj|up_proj|down_proj)\.weight$"
+)
+
+
+def _apply_buffer_loading_patch():
+    """Patch the DeepSpeed ZeRO-3 state_dict loader to handle:
+    1. Key renaming (inner checkpoint format -> model format)
+    2. Per-expert -> 3D fused expert tensors
+    3. Manual buffer loading (e_score_correction_bias etc.)
+
+    ZeRO-3's _load_state_dict_into_zero3_model only handles named_parameters.
+    Buffers like e_score_correction_bias must be loaded manually.
+    """
+    try:
+        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
+    except ImportError:
+        logger.warning(
+            "Could not import transformers.integrations.deepspeed; "
+            "buffer loading patch NOT applied."
+        )
+        return
+
+    def _patched_load_zero3(model_to_load, state_dict, *args, **kwargs):
+        # Step 1: Key rename + per-expert collection
+        new_sd = {}
+        expert_groups = {}  # prefix -> {expert_idx -> {proj_name -> tensor}}
+
+        for k, v in state_dict.items():
+            m = _EXPERT_KEY_RE.match(k)
+            if m:
+                # Per-expert key: collect for fusion
+                prefix = m.group(1)
+                expert_idx = int(m.group(2))
+                proj_name = m.group(3)
+                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] = v
+            else:
+                # Non-expert key: apply simple renames
+                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
+
+        # Step 2: Fuse expert groups into 3D tensors
+        if expert_groups:
+            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:
+                        logger.warning(
+                            "HYV3 Patch 1: Missing expert %d in %s", i, prefix
+                        )
+                        continue
+                    exp = experts[i]
+                    if "gate_proj" in exp and "up_proj" in exp:
+                        gate_up_list.append(
+                            torch.cat([exp["gate_proj"], exp["up_proj"]], dim=0)
+                        )
+                    if "down_proj" in exp:
+                        down_list.append(exp["down_proj"])
+                if gate_up_list:
+                    new_sd[f"{prefix}gate_up_proj"] = torch.stack(gate_up_list, dim=0)
+                if down_list:
+                    new_sd[f"{prefix}down_proj"] = torch.stack(down_list, dim=0)
+            logger.info(
+                "HYV3 Patch 1: Fused %d expert groups from per-expert to 3D format.",
+                len(expert_groups)
+            )
+            del expert_groups
+
+        # Step 3: Load parameters via original ZeRO-3 loader
+        result = _orig_load_zero3(model_to_load, new_sd, *args, **kwargs)
+
+        # Step 4: 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
+                    if isinstance(result, tuple) and len(result) >= 2:
+                        if isinstance(result[1], set):
+                            result[1].discard(name)
+        if buffers_loaded > 0:
+            logger.info(
+                "HYV3 Patch 1: Manually loaded %d buffers into model.",
+                buffers_loaded
+            )
+
+        del new_sd
+        return result
+
+    _ds_mod._load_state_dict_into_zero3_model = _patched_load_zero3
+    _mu_mod._load_state_dict_into_zero3_model = _patched_load_zero3
+    logger.info(
+        "HYV3 patch applied: ZeRO-3 key rename + expert fuse + buffer loading."
+    )
+
+# ============================================================================
+# Patch 2: 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 patch 2 (tokenizer copy) 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_save(self, args, state, control, **kwargs):
+            # --- Patch 2: 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."
+    )
+
+# ============================================================================
+# Patch 3: Memory-efficient shard-by-shard model loading for ZeRO-3
+#
+# The default transformers from_pretrained + ZeRO-3 path loads ALL shards
+# into a single state_dict in CPU memory before distributing. For a ~670GB
+# model with 8 processes per node, this causes CPU OOM.
+#
+# This patch replaces from_pretrained with a shard-by-shard loader that:
+#   1. Creates the model skeleton under deepspeed.zero.Init (meta tensors)
+#   2. Loads each safetensors shard one at a time (~7GB each)
+#   3. Applies key renames + expert fusion per shard
+#   4. Scatters into ZeRO-3 partitions immediately
+#   5. Frees the shard before loading the next one
+#
+# This reduces per-rank CPU memory from ~670GB to ~7GB.
+# ============================================================================
+
+def _apply_shard_loading_patch():
+    """Monkey-patch AutoModelForCausalLM.from_pretrained to use shard-by-shard
+    loading when DeepSpeed ZeRO-3 is active."""
+    import gc
+    import json as _json
+    import transformers
+
+    _orig_from_pretrained = transformers.AutoModelForCausalLM.from_pretrained
+
+    def _shard_loading_from_pretrained(pretrained_model_name_or_path, *args, **kwargs):
+        """Memory-efficient from_pretrained that loads shards one at a time."""
+        import deepspeed
+
+        model_path = pretrained_model_name_or_path
+
+        # Only apply shard loading if:
+        # 1. It's a local directory with safetensors
+        # 2. DeepSpeed ZeRO-3 is being used
+        if not (isinstance(model_path, str) and os.path.isdir(model_path)):
+            return _orig_from_pretrained(pretrained_model_name_or_path, *args, **kwargs)
+
+        index_file = os.path.join(model_path, "model.safetensors.index.json")
+        single_file = os.path.join(model_path, "model.safetensors")
+        if not (os.path.isfile(index_file) or os.path.isfile(single_file)):
+            return _orig_from_pretrained(pretrained_model_name_or_path, *args, **kwargs)
+
+        # Try to determine the DeepSpeed config
+        ds_config = None
+
+        # Check if there's a deepspeed config in the HfTrainerDeepSpeedConfig
+        try:
+            from transformers.integrations.deepspeed import is_deepspeed_zero3_enabled
+            if not is_deepspeed_zero3_enabled():
+                logger.info(
+                    "[HYV3 Patch 3] ZeRO-3 not enabled, using default loader."
+                )
+                return _orig_from_pretrained(pretrained_model_name_or_path, *args, **kwargs)
+        except (ImportError, Exception):
+            # If we can't determine, try to proceed anyway
+            pass
+
+        # Get the deepspeed config from HF's global state
+        try:
+            from transformers.integrations.deepspeed import deepspeed_config as _get_ds_config
+            ds_config = _get_ds_config()
+        except (ImportError, Exception):
+            ds_config = None
+
+        if ds_config is None:
+            # Fallback: try weak ref approach
+            try:
+                from transformers.integrations import deepspeed as _hf_ds
+                if hasattr(_hf_ds, '_hf_deepspeed_config_weak_ref'):
+                    _weak_ref = _hf_ds._hf_deepspeed_config_weak_ref
+                    if _weak_ref is not None:
+                        ds_obj = _weak_ref()
+                        if ds_obj is not None:
+                            ds_config = ds_obj.config
+            except (ImportError, AttributeError, Exception):
+                pass
+
+        if ds_config is None:
+            # Last resort: look for the config file path in environment
+            ds_config_path = os.environ.get("DEEPSPEED_CONFIG_FILE", None)
+            if ds_config_path is None:
+                ds_config_path = os.environ.get("DEEPSPEED_CONFIG", None)
+            if ds_config_path and os.path.isfile(ds_config_path):
+                with open(ds_config_path, "r") as f:
+                    ds_config = _json.load(f)
+
+        if ds_config is None:
+            logger.warning(
+                "[HYV3 Patch 3] Cannot determine DeepSpeed config, "
+                "falling back to default from_pretrained."
+            )
+            return _orig_from_pretrained(pretrained_model_name_or_path, *args, **kwargs)
+
+        # Ensure ds_config is a dict
+        if hasattr(ds_config, 'config'):
+            ds_config = ds_config.config
+        if not isinstance(ds_config, dict):
+            logger.warning(
+                "[HYV3 Patch 3] ds_config is not a dict (%s), falling back.",
+                type(ds_config)
+            )
+            return _orig_from_pretrained(pretrained_model_name_or_path, *args, **kwargs)
+
+        # Check if it's actually ZeRO stage 3
+        zero_stage = ds_config.get("zero_optimization", {}).get("stage", 0)
+        if zero_stage != 3:
+            logger.info(
+                "[HYV3 Patch 3] Not ZeRO-3 (stage=%d), using default loader.",
+                zero_stage
+            )
+            return _orig_from_pretrained(pretrained_model_name_or_path, *args, **kwargs)
+
+        logger.info(
+            "[HYV3 Patch 3] Using shard-by-shard loading for model at: %s",
+            model_path
+        )
+
+        try:
+            from safetensors import safe_open
+            from transformers.integrations.deepspeed import (
+                _load_state_dict_into_zero3_model as _load_zero3,
+            )
+        except ImportError as e:
+            logger.warning(
+                "[HYV3 Patch 3] Required imports not available (%s), "
+                "falling back to default from_pretrained.", e
+            )
+            return _orig_from_pretrained(pretrained_model_name_or_path, *args, **kwargs)
+
+        # Replace "auto" values that deepspeed.zero.Init cannot resolve
+        ds_config_copy = _json.loads(_json.dumps(ds_config))
+        _auto_defaults = {
+            "train_batch_size": 32,
+            "train_micro_batch_size_per_gpu": 1,
+            "gradient_accumulation_steps": 1,
+            "gradient_clipping": 1.0,
+        }
+        for k, v in _auto_defaults.items():
+            if k in ds_config_copy and ds_config_copy[k] == "auto":
+                ds_config_copy[k] = v
+
+        # Determine dtype
+        torch_dtype = kwargs.pop("torch_dtype", torch.bfloat16)
+        if torch_dtype is None or torch_dtype == "auto":
+            torch_dtype = torch.bfloat16
+        trust_remote_code = kwargs.pop("trust_remote_code", True)
+        attn_implementation = kwargs.pop("attn_implementation", None)
+        # Pop config if already provided by caller (e.g. LLaMA Factory)
+        config = kwargs.pop("config", None)
+
+        # Step 1: Create model skeleton under ZeRO-3 Init (meta tensors)
+        if config is None:
+            config = transformers.AutoConfig.from_pretrained(
+                model_path, trust_remote_code=trust_remote_code
+            )
+        with deepspeed.zero.Init(
+            dtype=torch_dtype, config_dict_or_path=ds_config_copy
+        ):
+            model = transformers.AutoModelForCausalLM.from_config(
+                config,
+                trust_remote_code=trust_remote_code,
+                torch_dtype=torch_dtype,
+                attn_implementation=attn_implementation,
+            )
+        logger.info("[HYV3 Patch 3] Model skeleton created under ZeRO-3 Init.")
+
+        # Step 2: Determine shard files
+        if os.path.isfile(index_file):
+            with open(index_file, "r") as f:
+                index_data = _json.load(f)
+            shard_files = list(dict.fromkeys(index_data["weight_map"].values()))
+        else:
+            shard_files = ["model.safetensors"]
+
+        # Step 3: Load each shard and scatter into ZeRO-3 model
+        total_shards = len(shard_files)
+        all_loaded_keys = set()
+        pending_experts = {}  # prefix -> {expert_idx -> {proj_name -> tensor}}
+
+        for shard_idx, shard_name in enumerate(shard_files, 1):
+            shard_path = os.path.join(model_path, shard_name)
+            logger.info(
+                "[HYV3 Patch 3] Loading shard %d/%d: %s",
+                shard_idx, total_shards, shard_name
+            )
+
+            # Load shard into CPU memory
+            shard_sd = {}
+            with safe_open(shard_path, framework="pt", device="cpu") as f:
+                for key in f.keys():
+                    shard_sd[key] = f.get_tensor(key)
+
+            # Separate expert keys from non-expert keys, apply renames
+            renamed_sd = {}
+            expert_keys_in_shard = {}
+
+            for k, v in shard_sd.items():
+                m = _EXPERT_KEY_RE.match(k)
+                if m:
+                    prefix = m.group(1)
+                    expert_idx = int(m.group(2))
+                    proj_name = m.group(3)
+                    if prefix not in expert_keys_in_shard:
+                        expert_keys_in_shard[prefix] = {}
+                    if expert_idx not in expert_keys_in_shard[prefix]:
+                        expert_keys_in_shard[prefix][expert_idx] = {}
+                    expert_keys_in_shard[prefix][expert_idx][proj_name] = v
+                else:
+                    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
+                    renamed_sd[new_k] = v
+            del shard_sd
+
+            # Merge expert keys into pending_experts
+            for prefix, experts in expert_keys_in_shard.items():
+                if prefix not in pending_experts:
+                    pending_experts[prefix] = {}
+                for idx, projs in experts.items():
+                    if idx not in pending_experts[prefix]:
+                        pending_experts[prefix][idx] = {}
+                    pending_experts[prefix][idx].update(projs)
+            del expert_keys_in_shard
+
+            # Check for completed expert groups
+            completed_prefixes = []
+            for prefix, experts in pending_experts.items():
+                if not experts:
+                    continue
+                max_idx = max(experts.keys())
+                num_experts_found = len(experts)
+                all_complete = all(
+                    len(projs) == 3 for projs in experts.values()
+                )
+                if all_complete and num_experts_found == (max_idx + 1):
+                    completed_prefixes.append(prefix)
+
+            # Fuse completed expert groups
+            for prefix in completed_prefixes:
+                experts = pending_experts.pop(prefix)
+                num_experts_layer = max(experts.keys()) + 1
+                gate_up_list = []
+                down_list = []
+                for i in range(num_experts_layer):
+                    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)
+                del gate_up_list, down_list, experts
+                renamed_sd[f"{prefix}gate_up_proj"] = fused_gate_up
+                renamed_sd[f"{prefix}down_proj"] = fused_down
+                logger.info(
+                    "[HYV3 Patch 3]   Fused %d experts for %s",
+                    num_experts_layer, prefix
+                )
+
+            # Scatter this shard's weights into ZeRO-3 model
+            if renamed_sd:
+                _load_zero3(model, renamed_sd)
+                # Also load buffers
+                for name, buf in model.named_buffers():
+                    if name in renamed_sd:
+                        src_tensor = renamed_sd[name]
+                        if isinstance(src_tensor, torch.Tensor):
+                            buf.data.copy_(src_tensor.to(buf.dtype))
+                all_loaded_keys.update(renamed_sd.keys())
+            del renamed_sd
+            gc.collect()
+
+        # Flush remaining pending experts
+        if pending_experts:
+            logger.info(
+                "[HYV3 Patch 3] Flushing %d remaining expert group(s)...",
+                len(pending_experts)
+            )
+            flush_sd = {}
+            for prefix, experts in pending_experts.items():
+                num_experts_layer = max(experts.keys()) + 1
+                gate_up_list = []
+                down_list = []
+                for i in range(num_experts_layer):
+                    if i not in experts:
+                        logger.warning(
+                            "[HYV3 Patch 3] Missing expert %d in %s", i, prefix
+                        )
+                        continue
+                    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"])
+                if gate_up_list:
+                    fused_gate_up = torch.stack(gate_up_list, dim=0)
+                    fused_down = torch.stack(down_list, dim=0)
+                    flush_sd[f"{prefix}gate_up_proj"] = fused_gate_up
+                    flush_sd[f"{prefix}down_proj"] = fused_down
+                    logger.info(
+                        "[HYV3 Patch 3]   Fused %d experts for %s",
+                        len(gate_up_list), prefix
+                    )
+                del gate_up_list, down_list
+            del pending_experts
+
+            if flush_sd:
+                _load_zero3(model, flush_sd)
+                for name, buf in model.named_buffers():
+                    if name in flush_sd:
+                        src_tensor = flush_sd[name]
+                        if isinstance(src_tensor, torch.Tensor):
+                            buf.data.copy_(src_tensor.to(buf.dtype))
+                all_loaded_keys.update(flush_sd.keys())
+            del flush_sd
+            gc.collect()
+
+        # Report missing/unexpected keys
+        model_keys = set(n for n, _ in model.named_parameters())
+        model_keys.update(n for n, _ in model.named_buffers())
+        missing = model_keys - all_loaded_keys
+        unexpected = all_loaded_keys - model_keys
+        if missing:
+            real_missing = {k for k in missing if "lm_head" not in k}
+            if real_missing:
+                logger.warning(
+                    "[HYV3 Patch 3] %d keys not found in checkpoint (first 10): %s",
+                    len(real_missing), list(real_missing)[:10]
+                )
+        if unexpected:
+            logger.warning(
+                "[HYV3 Patch 3] %d unexpected keys (first 10): %s",
+                len(unexpected), list(unexpected)[:10]
+            )
+        logger.info(
+            "[HYV3 Patch 3] Shard-by-shard loading complete. "
+            "Loaded %d keys from %d shards.",
+            len(all_loaded_keys), total_shards
+        )
+
+        return model
+
+    # Apply the monkey-patch
+    transformers.AutoModelForCausalLM.from_pretrained = staticmethod(_shard_loading_from_pretrained)
+    logger.info(
+        "HYV3 patch applied: shard-by-shard model loading for ZeRO-3 "
+        "(reduces CPU memory from ~670GB to ~7GB per rank)."
+    )
+
+
+# ============================================================================
+# Auto-apply patches on import
+# ============================================================================
+
+# Patch 1: ZeRO-3 key rename + expert fuse + buffer loading
+_apply_buffer_loading_patch()
+
+# Patch 3: Memory-efficient shard-by-shard loading
+_apply_shard_loading_patch()
+
+# Patch 2 (tokenizer copy) is applied via HYV3PatchCallback during training.
+# Users should add HYV3PatchCallback to their Trainer callbacks.
+
+logger.info(
+    "HYV3 patches module loaded. Remember to add HYV3PatchCallback to "
+    "your Trainer callbacks for tokenizer file copy on save."
+)

finetune/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,
+)

finetune/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()

finetune/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="${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}"

finetune/ms_swift_support/hy_v3_full_sft.yaml

@@ -0,0 +1,115 @@
+# ============================================================================
+# ms-swift SFT config for HYV3 Full-parameter Fine-tuning
+#
+# Usage:
+#   YAML_FILE=hy_v3_full_sft.yaml bash sft_train.sh
+#
+# or directly:
+#   swift sft --config hy_v3_full_sft.yaml
+#
+# Model: Tencent-Hunyuan/Hy3
+# Template: hy_v3 (natively supported in ms-swift 4.2.2)
+# ============================================================================
+
+# ---- Model Settings ----
+model: Tencent-Hunyuan/Hy3
+# model: /path/to/local/Hy3  # use local path if needed
+model_type: hy_v3
+template: hy_v3
+torch_dtype: bfloat16
+attn_impl: flash_attn
+# experts_impl: auto  # let ms-swift choose the best MoE expert implementation
+
+# ---- Training Method ----
+# Full-parameter SFT: no LoRA settings needed
+# For LoRA, uncomment below:
+# tuner_backend: peft
+# tuner_type: lora
+# lora_rank: 8
+# lora_alpha: 16
+# lora_dropout: 0.05
+# target_modules:
+#   - q_proj
+#   - k_proj
+#   - v_proj
+#   - o_proj
+#   - gate_proj
+#   - up_proj
+#   - down_proj
+
+# ---- Dataset Settings ----
+# Option 1: Use dataset registered in dataset_info.json
+# dataset:
+#   - hy_v3_demo
+# dataset_dir: .  # path to dataset_info.json
+
+# Option 2: Use local jsonl file directly (RECOMMENDED)
+# ms-swift supports sharegpt-format jsonl directly:
+dataset:
+  - ../data/example_data.jsonl
+
+max_length: 4096  # HYV3 supports up to 262k context
+truncation_strategy: delete  # or "truncation_left"
+lazy_tokenize: true
+preprocess_num_proc: 4
+
+# ---- Output Settings ----
+output_dir: saves/hy_v3/full/sft
+save_steps: 500
+save_strategy: steps
+save_total_limit: 3
+save_only_model: false
+logging_steps: 10
+report_to: none  # choices: [none, wandb, tensorboard, swanlab, mlflow]
+
+# ---- Training Hyperparameters ----
+per_device_train_batch_size: 1
+gradient_accumulation_steps: 8
+learning_rate: 1.0e-5
+num_train_epochs: 3.0
+max_steps: -1  # set to >0 to override num_train_epochs
+warmup_ratio: 0.1
+lr_scheduler_type: cosine
+# lr_scheduler_kwargs:
+#   min_lr_rate: 0.1  # min_lr / learning_rate
+bf16: true
+# fp16: false
+
+# ---- DeepSpeed / Optimization ----
+deepspeed: zero3_offload  # options: zero0, zero2, zero2_offload, zero3, zero3_offload
+# You can also specify a custom deepspeed config file:
+# deepspeed: /path/to/ds_config.json
+gradient_checkpointing: true
+# gradient_checkpointing_kwargs:
+#   use_reentrant: true  # may be needed for ZeRO-3, test first
+max_grad_norm: 1.0
+weight_decay: 0.1
+adam_beta1: 0.9
+adam_beta2: 0.95
+optim: adamw_torch
+
+# ---- Distributed Training ----
+ddp_timeout: 180000000
+# packing: true  # enable sequence packing for efficiency
+# packing_length: 4096
+
+# ---- Evaluation (optional) ----
+# eval_strategy: steps
+# eval_steps: 500
+# eval_dataset:
+#   - hy_v3_eval
+# per_device_eval_batch_size: 1
+# eval_use_evalscope: false
+
+# ---- Generation Settings (for inference during/after training) ----
+max_new_tokens: 2048
+temperature: 0.7
+top_p: 0.9
+# stop_words:
+#   - "</s>"
+#   - "<|im_end|>"
+
+# ---- Misc ----
+seed: 42
+ignore_data_skip: true
+resume_from_checkpoint: null  # set to checkpoint path to resume

finetune/ms_swift_support/hy_v3_lora_sft.yaml

@@ -0,0 +1,104 @@
+# ============================================================================
+# ms-swift SFT config for HYV3 LoRA Fine-tuning
+#
+# Usage:
+#   YAML_FILE=hy_v3_lora_sft.yaml bash sft_train.sh
+#
+# or directly:
+#   swift sft --config hy_v3_lora_sft.yaml
+#
+# Model: Tencent-Hunyuan/Hy3
+# Template: hy_v3 (natively supported in ms-swift 4.2.2)
+# Training method: LoRA (Low-Rank Adaptation)
+# ============================================================================
+
+# ---- Model Settings ----
+model: Tencent-Hunyuan/Hy3
+# model: /path/to/local/Hy3  # use local path if needed
+model_type: hy_v3
+template: hy_v3
+torch_dtype: bfloat16
+attn_impl: flash_attn
+# experts_impl: auto
+
+# ---- Training Method: LoRA ----
+tuner_backend: peft
+tuner_type: lora
+lora_rank: 8
+lora_alpha: 16
+lora_dropout: 0.05
+# target_modules:  # auto-detected if not specified
+#   - q_proj
+#   - k_proj
+#   - v_proj
+#   - o_proj
+#   - gate_proj
+#   - up_proj
+#   - down_proj
+# use_rslora: false
+# use_dora: false
+
+# ---- Dataset Settings ----
+# Option 1: Use dataset registered in dataset_info.json
+# dataset:
+#   - hy_v3_demo
+# dataset_dir: .
+
+# Option 2: Use local jsonl file directly (RECOMMENDED)
+dataset:
+  - ../data/example_data.jsonl
+
+max_length: 4096
+truncation_strategy: delete
+lazy_tokenize: true
+preprocess_num_proc: 4
+
+# ---- Output Settings ----
+output_dir: saves/hy_v3/lora/sft
+save_steps: 500
+save_strategy: steps
+save_total_limit: 3
+save_only_model: false
+logging_steps: 10
+report_to: none
+
+# ---- Training Hyperparameters ----
+per_device_train_batch_size: 1
+gradient_accumulation_steps: 8
+learning_rate: 3.0e-4  # LoRA typically uses higher LR
+num_train_epochs: 3.0
+max_steps: -1
+warmup_ratio: 0.1
+lr_scheduler_type: cosine
+bf16: true
+
+# ---- DeepSpeed / Optimization ----
+deepspeed: ../deepspeed_support/ds_zero3_offload.json  # Use custom ZeRO-3 offload config
+gradient_checkpointing: true
+max_grad_norm: 1.0
+weight_decay: 0.1
+adam_beta1: 0.9
+adam_beta2: 0.95
+optim: adamw_torch
+
+# ---- Distributed Training ----
+ddp_timeout: 180000000
+# packing: true
+# packing_length: 4096
+
+# ---- Evaluation (optional) ----
+# eval_strategy: steps
+# eval_steps: 500
+# eval_dataset:
+#   - hy_v3_eval
+# per_device_eval_batch_size: 1
+
+# ---- Generation Settings ----
+max_new_tokens: 2048
+temperature: 0.7
+top_p: 0.9
+
+# ---- Misc ----
+seed: 42
+ignore_data_skip: true
+resume_from_checkpoint: null

finetune/ms_swift_support/hy_v3_swift_patches.py

@@ -0,0 +1,309 @@
+"""
+HYV3 patches for ms-swift + DeepSpeed ZeRO-3 training.
+
+This module applies necessary runtime patches so that HYV3 (MoE) can be
+trained correctly under ms-swift with DeepSpeed ZeRO-3.
+
+Patches applied:
+    1. Template fix: Re-register hy_v3 template with dynamic [['eos_token_id']]
+       for chat_sep and suffix (fixes inference stop token issue).
+    2. Shard-by-shard model loading (Patch 3): Replaces the default
+       from_pretrained which loads ALL shards into CPU memory at once,
+       causing OOM for large models (~670GB). Instead, loads one shard
+       at a time (~7GB each), leveraging transformers 5.8.1's built-in
+       conversion_mapping for key rename + expert fusion.
+
+Usage:
+    swift sft --custom_register_path hy_v3_swift_patches.py ...
+"""
+
+import os
+import gc
+import json as _json
+import logging
+from typing import Optional
+
+import torch
+
+logger = logging.getLogger(__name__)
+
+# ============================================================================
+# Patch 1: Template fix
+# ============================================================================
+
+from swift.template.register import TemplateMeta, register_template
+from swift.template.constant import LLMTemplateType
+from swift.template.templates.llm import HyV3Template
+
+register_template(
+    TemplateMeta(
+        LLMTemplateType.hy_v3,
+        prefix=['<｜hy_begin▁of▁sentence｜>'],
+        system_prefix=['<｜hy_begin▁of▁sentence｜>{{SYSTEM}}'],
+        prompt=['<｜hy_User｜>{{QUERY}}<｜hy_Assistant｜>'],
+        chat_sep=[['eos_token_id']],
+        suffix=[['eos_token_id']],
+        template_cls=HyV3Template,
+        is_thinking=True,
+        thinking_prefix='',
+        non_thinking_prefix='',
+        history_thinking_prefix='',
+        agent_template='hy_v3',
+    ),
+    exist_ok=True,
+)
+
+logger.info(
+    "HYV3 template patch applied: hy_v3 template re-registered with "
+    "dynamic [['eos_token_id']] for chat_sep and suffix."
+)
+
+# ============================================================================
+# Patch 3: Memory-efficient shard-by-shard model loading for ZeRO-3
+#
+# The default transformers 5.8.1 from_pretrained + ZeRO-3 path loads ALL
+# shards into a single merged_state_dict in CPU memory before distributing.
+# For a ~670GB model with 8 processes per node, this causes CPU OOM.
+#
+# This patch replaces from_pretrained with a shard-by-shard loader that:
+#   1. Creates the model skeleton under deepspeed.zero.Init (meta tensors)
+#   2. Loads each safetensors shard one at a time (~7GB each)
+#   3. Passes each shard through _load_state_dict_into_zero3_model which
+#      internally applies the conversion_mapping (key rename + expert fusion)
+#   4. Frees the shard before loading the next one
+#
+# This reduces per-rank CPU memory from ~670GB to ~7GB.
+#
+# Note: Unlike the LLaMA-Factory version, we do NOT need to manually handle
+# key renames or expert fusion here, because transformers 5.8.1's
+# _load_state_dict_into_zero3_model already applies weight_mapping
+# (conversion_mapping) internally.
+# ============================================================================
+
+def _apply_shard_loading_patch():
+    """Monkey-patch AutoModelForCausalLM.from_pretrained to use shard-by-shard
+    loading when DeepSpeed ZeRO-3 is active."""
+    import transformers
+
+    _orig_from_pretrained = transformers.AutoModelForCausalLM.from_pretrained
+
+    def _shard_loading_from_pretrained(pretrained_model_name_or_path, *args, **kwargs):
+        """Memory-efficient from_pretrained that loads shards one at a time."""
+        import deepspeed
+
+        model_path = pretrained_model_name_or_path
+
+        # Only apply shard loading if:
+        # 1. It's a local directory with safetensors
+        # 2. DeepSpeed ZeRO-3 is being used
+        if not (isinstance(model_path, str) and os.path.isdir(model_path)):
+            return _orig_from_pretrained(pretrained_model_name_or_path, *args, **kwargs)
+
+        index_file = os.path.join(model_path, "model.safetensors.index.json")
+        single_file = os.path.join(model_path, "model.safetensors")
+        if not (os.path.isfile(index_file) or os.path.isfile(single_file)):
+            return _orig_from_pretrained(pretrained_model_name_or_path, *args, **kwargs)
+
+        # Check if ZeRO-3 is enabled
+        try:
+            from transformers.integrations.deepspeed import is_deepspeed_zero3_enabled
+            if not is_deepspeed_zero3_enabled():
+                logger.info(
+                    "[HYV3 Patch 3] ZeRO-3 not enabled, using default loader."
+                )
+                return _orig_from_pretrained(pretrained_model_name_or_path, *args, **kwargs)
+        except (ImportError, Exception):
+            pass
+
+        # Get the deepspeed config
+        ds_config = None
+        try:
+            from transformers.integrations.deepspeed import deepspeed_config as _get_ds_config
+            ds_config = _get_ds_config()
+        except (ImportError, Exception):
+            ds_config = None
+
+        if ds_config is None:
+            try:
+                from transformers.integrations import deepspeed as _hf_ds
+                if hasattr(_hf_ds, '_hf_deepspeed_config_weak_ref'):
+                    _weak_ref = _hf_ds._hf_deepspeed_config_weak_ref
+                    if _weak_ref is not None:
+                        ds_obj = _weak_ref()
+                        if ds_obj is not None:
+                            ds_config = ds_obj.config
+            except (ImportError, AttributeError, Exception):
+                pass
+
+        if ds_config is None:
+            ds_config_path = os.environ.get("DEEPSPEED_CONFIG_FILE", None)
+            if ds_config_path is None:
+                ds_config_path = os.environ.get("DEEPSPEED_CONFIG", None)
+            if ds_config_path and os.path.isfile(ds_config_path):
+                with open(ds_config_path, "r") as f:
+                    ds_config = _json.load(f)
+
+        if ds_config is None:
+            logger.warning(
+                "[HYV3 Patch 3] Cannot determine DeepSpeed config, "
+                "falling back to default from_pretrained."
+            )
+            return _orig_from_pretrained(pretrained_model_name_or_path, *args, **kwargs)
+
+        # Ensure ds_config is a dict
+        if hasattr(ds_config, 'config'):
+            ds_config = ds_config.config
+        if not isinstance(ds_config, dict):
+            logger.warning(
+                "[HYV3 Patch 3] ds_config is not a dict (%s), falling back.",
+                type(ds_config)
+            )
+            return _orig_from_pretrained(pretrained_model_name_or_path, *args, **kwargs)
+
+        # Check if it's actually ZeRO stage 3
+        zero_stage = ds_config.get("zero_optimization", {}).get("stage", 0)
+        if zero_stage != 3:
+            logger.info(
+                "[HYV3 Patch 3] Not ZeRO-3 (stage=%d), using default loader.",
+                zero_stage
+            )
+            return _orig_from_pretrained(pretrained_model_name_or_path, *args, **kwargs)
+
+        logger.info(
+            "[HYV3 Patch 3] Using shard-by-shard loading for model at: %s",
+            model_path
+        )
+
+        try:
+            from safetensors import safe_open
+            from transformers.integrations.deepspeed import (
+                _load_state_dict_into_zero3_model as _load_zero3,
+            )
+            from transformers.conversion_mapping import get_model_conversion_mapping
+            from transformers.modeling_utils import LoadStateDictConfig
+        except ImportError as e:
+            logger.warning(
+                "[HYV3 Patch 3] Required imports not available (%s), "
+                "falling back to default from_pretrained.", e
+            )
+            return _orig_from_pretrained(pretrained_model_name_or_path, *args, **kwargs)
+
+        # Replace "auto" values that deepspeed.zero.Init cannot resolve
+        ds_config_copy = _json.loads(_json.dumps(ds_config))
+        _auto_defaults = {
+            "train_batch_size": 32,
+            "train_micro_batch_size_per_gpu": 1,
+            "gradient_accumulation_steps": 1,
+            "gradient_clipping": 1.0,
+        }
+        for k, v in _auto_defaults.items():
+            if k in ds_config_copy and ds_config_copy[k] == "auto":
+                ds_config_copy[k] = v
+
+        # Determine dtype - handle both torch_dtype (old) and dtype (new, transformers >= 4.56)
+        torch_dtype = kwargs.pop("torch_dtype", None)
+        if torch_dtype is None:
+            torch_dtype = kwargs.pop("dtype", torch.bfloat16)
+        if torch_dtype is None or torch_dtype == "auto":
+            torch_dtype = torch.bfloat16
+        if isinstance(torch_dtype, str):
+            torch_dtype = getattr(torch, torch_dtype, torch.bfloat16)
+
+        trust_remote_code = kwargs.pop("trust_remote_code", True)
+        attn_implementation = kwargs.pop("attn_implementation", None)
+        config = kwargs.pop("config", None)
+
+        # Step 1: Create model skeleton under ZeRO-3 Init (meta tensors)
+        if config is None:
+            config = transformers.AutoConfig.from_pretrained(
+                model_path, trust_remote_code=trust_remote_code
+            )
+
+        with deepspeed.zero.Init(
+            dtype=torch_dtype, config_dict_or_path=ds_config_copy
+        ):
+            model = transformers.AutoModelForCausalLM.from_config(
+                config,
+                trust_remote_code=trust_remote_code,
+                torch_dtype=torch_dtype,
+                attn_implementation=attn_implementation,
+            )
+        logger.info("[HYV3 Patch 3] Model skeleton created under ZeRO-3 Init.")
+
+        # Step 2: Get weight conversion mapping (key rename + expert fusion)
+        # transformers 5.8.1 has built-in conversion_mapping for hy_v3
+        weight_conversions = get_model_conversion_mapping(model, None, None)
+
+        # Create a minimal load_config with weight_mapping
+        load_config = LoadStateDictConfig(
+            pretrained_model_name_or_path=model_path,
+            weight_mapping=weight_conversions,
+        )
+
+        # Step 3: Determine shard files
+        if os.path.isfile(index_file):
+            with open(index_file, "r") as f:
+                index_data = _json.load(f)
+            shard_files = list(dict.fromkeys(index_data["weight_map"].values()))
+        else:
+            shard_files = ["model.safetensors"]
+
+        # Step 4: Load each shard and scatter into ZeRO-3 model
+        total_shards = len(shard_files)
+
+        for shard_idx, shard_name in enumerate(shard_files, 1):
+            shard_path = os.path.join(model_path, shard_name)
+            logger.info(
+                "[HYV3 Patch 3] Loading shard %d/%d: %s",
+                shard_idx, total_shards, shard_name
+            )
+
+            # Load shard into CPU memory
+            shard_sd = {}
+            with safe_open(shard_path, framework="pt", device="cpu") as f:
+                for key in f.keys():
+                    shard_sd[key] = f.get_tensor(key)
+
+            # Use transformers' built-in ZeRO-3 loader which handles:
+            # - weight_mapping (key rename + expert fusion via conversion_mapping)
+            # - buffer loading
+            # - parameter scattering into ZeRO-3 partitions
+            _load_zero3(model, shard_sd, load_config)
+
+            del shard_sd
+            gc.collect()
+
+        logger.info(
+            "[HYV3 Patch 3] Shard-by-shard loading complete. "
+            "Loaded %d shards.", total_shards
+        )
+
+        # Patch G: Disable output_router_logits to save CPU memory during training.
+        # When output_router_logits=True, all 79 MoE layers accumulate router logits
+        # tensors throughout forward pass, causing significant memory growth under
+        # ZeRO-3 offload. Since router_aux_loss_coef=0.0 (no aux loss), these logits
+        # are not needed for training.
+        if hasattr(model, 'config') and getattr(model.config, 'output_router_logits', False):
+            model.config.output_router_logits = False
+            logger.info(
+                "[HYV3 Patch G] Disabled output_router_logits to reduce "
+                "CPU memory usage during ZeRO-3 offload training."
+            )
+
+        return model
+
+    # Apply the monkey-patch
+    transformers.AutoModelForCausalLM.from_pretrained = staticmethod(_shard_loading_from_pretrained)
+    logger.info(
+        "HYV3 Patch 3 applied: shard-by-shard model loading for ZeRO-3 "
+        "(reduces CPU memory from ~670GB to ~7GB per rank)."
+    )
+
+
+# ============================================================================
+# Auto-apply patches on import
+# ============================================================================
+
+_apply_shard_loading_patch()
+
+logger.info("HYV3 ms-swift patches loaded successfully.")

finetune/ms_swift_support/sft_train.sh

@@ -0,0 +1,205 @@
+#!/bin/bash
+# ============================================================================
+# ms-swift SFT training launch script for HYV3
+#
+# ms-swift 4.2.2 has native HYV3 support:
+#   - Model registered: LLMModelType.hy_v3
+#   - Template registered: TemplateType.hy_v3
+#   - Agent template: HyV3AgentTemplate
+#   - No monkey-patches needed for basic full-parameter or LoRA SFT.
+#
+# Usage:
+#   Single node:  bash sft_train.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 sft_train.sh
+#
+# Note: ms-swift does NOT support --config parameter.
+#       All parameters must be passed directly via command line.
+# ============================================================================
+
+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
+
+# -------------------- 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]}
+
+# -------------------- 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 ms-swift SFT 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 --------------------
+# ms-swift does NOT support --config parameter.
+# All parameters must be passed directly via command line.
+# For multi-node, we need to set the distributed env vars and let swift handle it.
+
+# Common SFT parameters from hy_v3_full_sft.yaml
+SFT_PARAMS=(
+    # ---- Model Settings ----
+    --model /path/to/Hy3
+    --model_type hy_v3
+    --template hy_v3
+    --torch_dtype bfloat16
+    --tuner_type full
+    --attn_impl flash_attn
+
+    # ---- Dataset Settings ----
+    --dataset ../data/example_data.jsonl
+    --max_length 4096
+    --truncation_strategy delete
+    --lazy_tokenize true
+    --dataset_num_proc 4
+
+    # ---- Output Settings ----
+    --output_dir saves/hy_v3/full/sft
+    --save_steps 500
+    --save_strategy steps
+    --save_total_limit 3
+    --save_only_model false
+    --logging_steps 10
+    --report_to none
+
+    # ---- Training Hyperparameters ----
+    --per_device_train_batch_size 1
+    --gradient_accumulation_steps 1
+    --learning_rate 1.0e-5
+    --num_train_epochs 3.0
+    --max_steps -1
+    --warmup_ratio 0.1
+    --lr_scheduler_type cosine
+    --bf16 true
+
+    # ---- DeepSpeed / Optimization ----
+    --deepspeed zero3_offload
+    --gradient_checkpointing true
+    --max_grad_norm 1.0
+    --weight_decay 0.1
+    --adam_beta1 0.9
+    --adam_beta2 0.95
+    --optim adamw_torch
+
+    # ---- Distributed Training ----
+    --ddp_timeout 180000000
+
+    # ---- Generation Settings ----
+    --max_new_tokens 2048
+    --temperature 0.7
+    --top_p 0.9
+
+    # ---- Misc ----
+    --seed 42
+    --ignore_data_skip true
+)
+
+if [ ${NODES} -eq 1 ]; then
+    # Single-node: use torchrun to ensure local_world_size is set correctly
+    # This avoids the DeepSpeed + device_map compatibility error
+    export NODE_RANK=0
+    export NNODES=1
+
+    # Add current directory to PYTHONPATH so hy_v3_swift_patches can be imported
+    export PYTHONPATH="${PYTHONPATH:+${PYTHONPATH}:}$(pwd)"
+
+    torchrun \
+        --nproc_per_node "${HOST_GPU_NUM}" \
+        --master_port "${MASTER_PORT}" \
+        -m swift.cli.sft \
+        --custom_register_path hy_v3_swift_patches.py \
+        "${SFT_PARAMS[@]}"
+else
+    # Multi-node: use torchrun
+    # Determine local node rank
+    LOCAL_IP=$(hostname -i 2>/dev/null || 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 NODE_RANK="${NODE_RANK}"
+    export NNODES="${NODES}"
+    export MASTER_ADDR="${MASTER_ADDR}"
+    export MASTER_PORT="${MASTER_PORT}"
+
+    # Add current directory to PYTHONPATH so hy_v3_swift_patches can be imported
+    export PYTHONPATH="${PYTHONPATH:+${PYTHONPATH}:}$(pwd)"
+
+    torchrun \
+        --nnodes "${NNODES}" \
+        --node_rank "${NODE_RANK}" \
+        --nproc_per_node "${HOST_GPU_NUM}" \
+        --master_addr "${MASTER_ADDR}" \
+        --master_port "${MASTER_PORT}" \
+        -m swift.cli.sft \
+        --custom_register_path hy_v3_swift_patches.py \
+        "${SFT_PARAMS[@]}"
+fi

finetune/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

finetune/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()

finetune/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()