apps/FastDeploy
1
0
Fork 0
mirror of https://github.com/PaddlePaddle/FastDeploy.git synced 2025-10-05 08:37:06 +08:00
Code Issues Actions 2 Packages Projects Releases Wiki Activity

polish code with new pre-commit rule (#2923)

Browse Source
This commit is contained in:
Zero Rains
2025-07-19 23:19:27 +08:00
committed by GitHub
parent b8676d71a8
commit 25698d56d1
424 changed files with 14307 additions and 13518 deletions
Download Patch File Download Diff File Expand all files Collapse all files

115
fastdeploy/rl/rollout_model.py
View File

@@ -18,19 +18,32 @@ from typing import Dict
import paddle
from paddle import nn
from paddleformers.utils.log import logger
from fastdeploy.config import FDConfig
from fastdeploy.model_executor.model_loader import ModelRegistry
from fastdeploy.model_executor.models.ernie4_5_moe import \
Ernie4_5_MoeForCausalLM, Ernie4_5_PretrainedModel
from fastdeploy.model_executor.models.ernie4_5_vl.ernie4_5_vl_moe import \
Ernie4_5_VLMoeForConditionalGeneration, Ernie4_5_VLPretrainedModel
from fastdeploy.model_executor.models.qwen2 import Qwen2ForCausalLM, Qwen2PretrainedModel
from fastdeploy.model_executor.models.qwen3 import Qwen3ForCausalLM, Qwen3PretrainedModel
from fastdeploy.model_executor.models.qwen3moe import Qwen3MoeForCausalLM, Qwen3MoePretrainedModel
from fastdeploy.model_executor.models.ernie4_5_moe import (
Ernie4_5_MoeForCausalLM,
Ernie4_5_PretrainedModel,
)
from fastdeploy.model_executor.models.ernie4_5_vl.ernie4_5_vl_moe import (
Ernie4_5_VLMoeForConditionalGeneration,
Ernie4_5_VLPretrainedModel,
)
from fastdeploy.model_executor.models.qwen2 import (
Qwen2ForCausalLM,
Qwen2PretrainedModel,
)
from fastdeploy.model_executor.models.qwen3 import (
Qwen3ForCausalLM,
Qwen3PretrainedModel,
)
from fastdeploy.model_executor.models.qwen3moe import (
Qwen3MoeForCausalLM,
Qwen3MoePretrainedModel,
)
from fastdeploy.rl.rollout_config import RolloutModelConfig
class RolloutModel(nn.Layer):
"""Main model class for rollout operations, supports multimodal components for train."""
@@ -53,7 +66,7 @@ class RolloutModel(nn.Layer):
def get_name_mappings_to_training(self) -> Dict[str, str]:
"""Get parameter name mappings between rollout and training models."""
return getattr(self.rollout_model, "get_name_mappings_to_training", lambda: {})()
def get_quantization_infer_keys(self) -> Dict[str, str]:
"""Get parameter name mappings between rollout and training models."""
return getattr(self.rollout_model, "get_quantization_infer_keys", lambda: {})()
@@ -66,7 +79,10 @@ class RolloutModel(nn.Layer):
class BaseRLModel(nn.Layer):
"""Base class for RL models with common functionality"""
def __init__(self,):
def __init__(
self,
):
super(BaseRLModel, self).__init__()
self.infer_to_train_mapping = {}
self.fd_config = None
@@ -74,15 +90,15 @@ class BaseRLModel(nn.Layer):
@classmethod
def name(cls) -> str:
return cls.__name__
def _update_base_mappings(self, base_name: str) -> None:
"""Common static mappings"""
static_mappings = {
f"{base_name}.embed_tokens.embeddings.weight": f"{base_name}.embed_tokens.weight",
"lm_head.linear.weight": "lm_head.weight"
"lm_head.linear.weight": "lm_head.weight",
}
self.infer_to_train_mapping.update(static_mappings)
def _complete_missing_mappings(self) -> None:
"""
Complete the mapping dictionary with keys that have identical names in inference and training.
@@ -94,12 +110,12 @@ class BaseRLModel(nn.Layer):
if getattr(self.fd_config.model_config, "tie_word_embeddings", False):
self.infer_to_train_mapping.pop("lm_head.linear.weight")
def get_quantization_infer_keys(self) -> list[str]:
"""Get quantization infer keys"""
quant_weight_key = []
if self.fd_config.quant_config.name() == "wint8":
""" RL only support weight_only_int8 now"""
"""RL only support weight_only_int8 now"""
for key in self.state_dict().keys():
if "scale" in key:
quant_weight_key.append(key.replace(".weight_scale", ".weight"))
@@ -107,10 +123,12 @@ class BaseRLModel(nn.Layer):
raise ValueError("Only 'wint8' quantization is supported in RL roullout.")
return quant_weight_key
class Ernie4_5_MoeForCausalLMRL(Ernie4_5_MoeForCausalLM, BaseRLModel):
"""
Ernie4_5_MoeForCausalLMRL
"""
_get_tensor_parallel_mappings = Ernie4_5_PretrainedModel._get_tensor_parallel_mappings
def __init__(self, fd_config: FDConfig):
@@ -134,15 +152,18 @@ class Ernie4_5_MoeForCausalLMRL(Ernie4_5_MoeForCausalLM, BaseRLModel):
self._update_base_mappings("ernie")
base_name = "ernie.layers"
# Helper function to add layer mappings
def _add_layer_mappings(layer_idx: int):
# MoE specific mappings
self.infer_to_train_mapping[f"{base_name}.{layer_idx}.mlp.fused_moe.gate_weight"] = \
f"{base_name}.{layer_idx}.mlp.gate.weight"
self.infer_to_train_mapping[
f"{base_name}.{layer_idx}.mlp.fused_moe.gate_weight"
] = f"{base_name}.{layer_idx}.mlp.gate.weight"
if self.fd_config.model_config.moe_use_aux_free:
self.infer_to_train_mapping[f"{base_name}.{layer_idx}.mlp.fused_moe.gate_correction_bias"] = \
f"{base_name}.{layer_idx}.mlp.moe_statics.e_score_correction_bias"
self.infer_to_train_mapping[
f"{base_name}.{layer_idx}.mlp.fused_moe.gate_correction_bias"
] = f"{base_name}.{layer_idx}.mlp.moe_statics.e_score_correction_bias"
# MoE experts mappings
for expert_idx in range(self.fd_config.model_config.moe_num_experts):
@@ -165,8 +186,10 @@ class Ernie4_5_MoeForCausalLMRL(Ernie4_5_MoeForCausalLM, BaseRLModel):
assert isinstance(self.fd_config.model_config.moe_layer_start_index, int)
# Process MoE layers
for layer_idx in range(self.fd_config.model_config.moe_layer_start_index,
self.fd_config.model_config.num_hidden_layers):
for layer_idx in range(
self.fd_config.model_config.moe_layer_start_index,
self.fd_config.model_config.num_hidden_layers,
):
_add_layer_mappings(layer_idx)
self._complete_missing_mappings()
@@ -178,6 +201,7 @@ class Ernie4_5_VLMoeForConditionalGenerationRL(Ernie4_5_VLMoeForConditionalGener
"""
Ernie4_5_VLMoeForConditionalGenerationRL
"""
_get_tensor_parallel_mappings = Ernie4_5_VLPretrainedModel._get_tensor_parallel_mappings
def __init__(self, fd_config: FDConfig):
@@ -206,25 +230,30 @@ class Ernie4_5_VLMoeForConditionalGenerationRL(Ernie4_5_VLMoeForConditionalGener
def _add_expert_mappings(layer_idx: int, moe_tag: str, expert_start: int):
# MoE specific mappings
gate_suffix = "" if moe_tag == "text" else "_1"
self.infer_to_train_mapping[f"{base_name}.{layer_idx}.mlp.{moe_tag}_fused_moe.gate_weight"] = \
f"{base_name}.{layer_idx}.mlp.gate.weight{gate_suffix}"
self.infer_to_train_mapping[
f"{base_name}.{layer_idx}.mlp.{moe_tag}_fused_moe.gate_weight"
] = f"{base_name}.{layer_idx}.mlp.gate.weight{gate_suffix}"
if self.fd_config.model_config.moe_use_aux_free:
self.infer_to_train_mapping[f"{base_name}.{layer_idx}.mlp.{moe_tag}_fused_moe.gate_correction_bias"] = \
f"{base_name}.{layer_idx}.mlp.moe_statics.e_score_correction_bias"
self.infer_to_train_mapping[
f"{base_name}.{layer_idx}.mlp.{moe_tag}_fused_moe.gate_correction_bias"
] = f"{base_name}.{layer_idx}.mlp.moe_statics.e_score_correction_bias"
# Initialize defaultdict for expert weights
from collections import defaultdict
from itertools import chain
def _generate_ranges(start, end, step=16, take=8):
"""生成 [start, start+take), [start+step, start+step+take), ... 直到 end"""
return chain(
*(range(i, min(i + take, end)) # 防止越界
for i in range(start, end, step)))
return chain(*(range(i, min(i + take, end)) for i in range(start, end, step))) # 防止越界
expert_mappings = defaultdict(list)
for expert_idx in _generate_ranges(expert_start, total_moe_num, expert_num_per_rank * 2, expert_num_per_rank):
for expert_idx in _generate_ranges(
expert_start,
total_moe_num,
expert_num_per_rank * 2,
expert_num_per_rank,
):
for ph in place_holders:
expert_mappings[f"{base_name}.{layer_idx}.mlp.{moe_tag}_fused_moe.up_gate_proj_weight"].append(
f"{base_name}.{layer_idx}.mlp.experts.{expert_idx}.up_gate_proj.{ph}"
@@ -273,6 +302,7 @@ class Qwen2ForCausalLMRL(Qwen2ForCausalLM, BaseRLModel):
"""
Qwen2ForCausalLMRL
"""
_get_tensor_parallel_mappings = Qwen2PretrainedModel._get_tensor_parallel_mappings
def __init__(self, fd_config: FDConfig):
@@ -295,15 +325,16 @@ class Qwen2ForCausalLMRL(Qwen2ForCausalLM, BaseRLModel):
# Initialize mapping dictionary
self._update_base_mappings("qwen2")
base_name = "qwen2.layers"
# Helper function to add layer mappings
def _add_layer_mappings(layer_idx):
# FFN mappings
for ph in place_holders:
self.infer_to_train_mapping[f"{base_name}.{layer_idx}.mlp.up_gate_proj.{ph}"] = \
f"{base_name}.{layer_idx}.mlp.gate_up_fused_proj.{ph}"
self.infer_to_train_mapping[
f"{base_name}.{layer_idx}.mlp.up_gate_proj.{ph}"
] = f"{base_name}.{layer_idx}.mlp.gate_up_fused_proj.{ph}"
for layer_idx in range(
self.fd_config.model_config.num_hidden_layers):
for layer_idx in range(self.fd_config.model_config.num_hidden_layers):
_add_layer_mappings(layer_idx)
self._complete_missing_mappings()
@@ -315,6 +346,7 @@ class Qwen3MoeForCausalLMRL(Qwen3MoeForCausalLM, BaseRLModel):
"""
Qwen3MoeForCausalLMRL
"""
_get_tensor_parallel_mappings = Qwen3MoePretrainedModel._get_tensor_parallel_mappings
def __init__(self, fd_config: FDConfig):
@@ -343,12 +375,14 @@ class Qwen3MoeForCausalLMRL(Qwen3MoeForCausalLM, BaseRLModel):
# Helper function to add layer mappings
def _add_layer_mappings(layer_idx: int):
# MoE specific mappings
self.infer_to_train_mapping[f"{base_name}.{layer_idx}.mlp.gate_weight"] = \
f"{base_name}.{layer_idx}.mlp.gate.weight"
self.infer_to_train_mapping[
f"{base_name}.{layer_idx}.mlp.gate_weight"
] = f"{base_name}.{layer_idx}.mlp.gate.weight"
if self.fd_config.moe_config.moe_use_aux_free:
self.infer_to_train_mapping[f"{base_name}.{layer_idx}.mlp.fused_moe.gate_correction_bias"] = \
f"{base_name}.{layer_idx}.mlp.moe_statics.e_score_correction_bias"
self.infer_to_train_mapping[
f"{base_name}.{layer_idx}.mlp.fused_moe.gate_correction_bias"
] = f"{base_name}.{layer_idx}.mlp.moe_statics.e_score_correction_bias"
# MoE experts mappings
for expert_idx in range(self.fd_config.moe_config.num_experts):
@@ -382,6 +416,7 @@ class Qwen3ForCausalLMRL(Qwen3ForCausalLM, BaseRLModel):
"""
Qwen3ForCausalLMRL
"""
_get_tensor_parallel_mappings = Qwen3PretrainedModel._get_tensor_parallel_mappings
def __init__(self, fd_config: FDConfig):
@@ -395,6 +430,6 @@ class Qwen3ForCausalLMRL(Qwen3ForCausalLM, BaseRLModel):
def name(self) -> str:
"""name"""
return "Qwen3ForCausalLMRL"
def get_name_mappings_to_training(self) -> Dict[str, str]:
pass