finetune/deepspeed_support/train.py

# 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()