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