From c07a951ee089004c77e3583f8d9210e4d02ed5d8 Mon Sep 17 00:00:00 2001 From: RangiLyu Date: Wed, 8 Jul 2026 19:12:40 +0800 Subject: [PATCH 1/2] [Feature] Support LoRA training and RL weight update --- tests/engine/test_dense_lora_train_engine.py | 213 ++++++++++ xtuner/v1/engine/train_engine.py | 5 + xtuner/v1/model/adapter/__init__.py | 0 xtuner/v1/model/adapter/lora.py | 382 ++++++++++++++++++ xtuner/v1/module/lora_linear/__init__.py | 0 .../module/lora_linear/lora_grouped_linear.py | 97 +++++ xtuner/v1/module/lora_linear/lora_linear.py | 100 +++++ xtuner/v1/rl/trainer/worker.py | 2 + xtuner/v1/rl/weight_update/update_weighter.py | 17 + xtuner/v1/rl/weight_update/weight_iterator.py | 9 + xtuner/v1/train/trainer.py | 8 + 11 files changed, 833 insertions(+) create mode 100644 tests/engine/test_dense_lora_train_engine.py create mode 100644 xtuner/v1/model/adapter/__init__.py create mode 100644 xtuner/v1/model/adapter/lora.py create mode 100644 xtuner/v1/module/lora_linear/__init__.py create mode 100644 xtuner/v1/module/lora_linear/lora_grouped_linear.py create mode 100644 xtuner/v1/module/lora_linear/lora_linear.py diff --git a/tests/engine/test_dense_lora_train_engine.py b/tests/engine/test_dense_lora_train_engine.py new file mode 100644 index 0000000000..80d23b612c --- /dev/null +++ b/tests/engine/test_dense_lora_train_engine.py @@ -0,0 +1,213 @@ +import os +import tempfile +import shutil +import time +import parametrize +import torch +import torch.distributed as dist +from xtuner._testing import DeterministicDDPTestCase +from transformers import AutoTokenizer + +from xtuner.v1.model.moe.moe import SequenceContext +from xtuner.v1.model.dense.qwen3 import Qwen3Dense8BConfig +from xtuner.v1.model.base import ModelItem +from xtuner.v1.loss.ce_loss import CELossConfig, CELossContextInputItem +from xtuner.v1.config import FSDPConfig, LRConfig, AdamWConfig +from xtuner.v1.engine.train_engine import TrainEngine +from torch.optim.lr_scheduler import LambdaLR +from xtuner.v1.utils import pad_to_max_length +from xtuner.v1.utils.device import get_device +from xtuner.v1.utils.test_utils import init_data_mesh +from xtuner.v1.model.adapter.lora import LoraConfig + + +# Qwen3 8B +QWEN3_PATH = os.environ["QWEN3_PATH"] +DEVICE = get_device() + + +class TestDenseEngine(DeterministicDDPTestCase): + @parametrize.parametrize( + "device,tp_size,sp_size", + [ + ("cuda", 1, 1), + ("cuda", 1, 2), + ], + ) + def test_dense_engine_train(self, device, tp_size, sp_size): + pg = self.create_pg(device) + + dense_cfg = Qwen3Dense8BConfig() + optim_cfg: AdamWConfig = AdamWConfig() + lr_cfg: LRConfig = LRConfig() + fsdp_cfg: FSDPConfig = FSDPConfig( + torch_compile=True, + cpu_offload=False, + tp_size=tp_size, + # hsdp_sharding_size=hsdp_sharding_size, + ) + + adapter_cfg = LoraConfig( + r=8, + lora_alpha=8, + lora_dropout=0, + target_modules=[ + "q_proj", + "k_proj", + "v_proj", + "o_proj", + "gate_proj", + "up_proj", + "down_proj", + ], + bias="none", + modules_to_save=["lm_head"], + ) + engine = TrainEngine( + model_cfg=dense_cfg, + optim_cfg=optim_cfg, + fsdp_cfg=fsdp_cfg, + adapter_cfg=adapter_cfg, + ) + engine.from_hf(hf_path=QWEN3_PATH) + + loss_cfg = CELossConfig() + + total_steps = 1000 + warmup_steps = total_steps * lr_cfg.warmup_ratio + + def warmup_fn(x): + return x / warmup_steps if x < warmup_steps else 1 + + lr_scheduler = LambdaLR(engine.optimizer, warmup_fn) + + tok = AutoTokenizer.from_pretrained(QWEN3_PATH) + txt = "根据国际地球自转和参考系服务机构的数据，今年夏天是自2020年以来第六次地球自转加速。7月9日将成为有史以来最短的一天，比平时短1.3到1.6毫秒。 " + input_ids = tok.encode(txt, return_tensors="pt").view(1, -1) + labels = input_ids.clone() + input_ids = input_ids[:, :-1] + labels = labels[:, 1:] + pack_len = 8192 - input_ids.shape[1] + input_ids = pad_to_max_length(input_ids, 0, max_length=8192) + labels = pad_to_max_length(labels, -100, max_length=8192) + losses = [] + + data_mesh = None + if sp_size > 1: + data_mesh = init_data_mesh(str(DEVICE), sp_size) + + for _ in range(10): + seq_ctx = SequenceContext.from_input_ids((input_ids,), device=DEVICE) + labels = labels.to(DEVICE) + seq_ctx.num_padding = pack_len + seq_ctx_list = [seq_ctx] + loss_ctx_input_list: list[CELossContextInputItem] = [ + CELossContextInputItem(shifted_labels=labels) + ] + LossContext = loss_cfg.loss_ctx_cls + batches_loss_kwargs = LossContext.build_batches_loss_kwargs( + loss_ctx_input_list, + loss_cfg, + ) + loss_kwargs = batches_loss_kwargs[0] + loss_ctx = LossContext(loss_cfg, loss_kwargs) + seq_ctx = seq_ctx_list[0] + engine_input = [ModelItem(seq_ctx=seq_ctx, loss_ctx=loss_ctx)] + loss_log, _ = engine.train_step(engine_input) + grad_norm = engine.clip_grad_norm() + engine.step_optimizer(grad_norm) + lr_scheduler.step() + losses.append(loss_log["reduced_llm_loss"]) + losses_ref = [2.57, 2.57, 2.57, 2.57, 2.57, 2.57, 2.56, 2.56, 2.54, 2.53] + for loss, loss_ref in zip(losses, losses_ref): + self.assertTrue( + abs(loss - loss_ref) < 0.02, + f"loss={loss}, loss_ref={loss_ref}, diff={abs(loss - loss_ref)}", + ) + + torch.cuda.empty_cache() + try: + dist.destroy_process_group(pg) + except: + pass + + @parametrize.parametrize( + "device,tp_size,hsdp_sharding_size", + [ + ("cuda", 1, 8), # todo: test ep8 and hsdp, OOM in 8 gpus + ], + ) + def test_save_and_load(self, device, tp_size, hsdp_sharding_size): + pg = self.create_pg(device) + + temp_dir = tempfile.mkdtemp() + if dist.get_rank() == 0: + temp_dir = [temp_dir] + else: + temp_dir = [None] + dist.broadcast_object_list(temp_dir, src=0) + temp_dir = temp_dir[0] + moe_cfg = Qwen3Dense8BConfig() + optim_cfg: AdamWConfig = AdamWConfig() + fsdp_cfg: FSDPConfig = FSDPConfig( + torch_compile=True, + cpu_offload=False, + tp_size=tp_size, + hsdp_sharding_size=hsdp_sharding_size, + ) + adapter_cfg = LoraConfig( + r=4, + lora_alpha=16, + lora_dropout=0, + target_modules=["q_proj", "v_proj"], + bias="none", + modules_to_save=["lm_head"], + ) + engine = TrainEngine( + model_cfg=moe_cfg, + optim_cfg=optim_cfg, + fsdp_cfg=fsdp_cfg, + adapter_cfg=adapter_cfg, + ) + + engine.from_hf(hf_path=QWEN3_PATH) + engine.save_hf( + hf_dir=temp_dir, + save_dtype=torch.bfloat16, + ) + + dist.barrier() + time.sleep(1) + + # engine2 = TrainEngine( + # model_cfg=moe_cfg, + # optim_cfg=optim_cfg, + # fsdp_cfg=fsdp_cfg, + # ) + # engine2.from_hf(hf_path=temp_dir) + + # state_dict = engine.model.state_dict() + # state_dict2 = engine2.model.state_dict() + # for key, val in state_dict.items(): + # val2 = state_dict2[key] + # val = val.full_tensor().bfloat16() + # val2 = val2.full_tensor().bfloat16() + # self.assertTrue(torch.equal(val, val2[:val.shape[0]]), + # f"Mismatch in {key} between bf16 and fp8, {val} and {val2[:val.shape[0]]}") + + if dist.get_rank() == 0: + shutil.rmtree(temp_dir) + + torch.cuda.empty_cache() + try: + dist.destroy_process_group(pg) + except: + pass + + @property + def world_size(self) -> int: + return int(os.getenv("XTUNER_TEST_WORLD_SIZE", "8")) + + @property + def destroy_pg_upon_exit(self) -> bool: + return False \ No newline at end of file diff --git a/xtuner/v1/engine/train_engine.py b/xtuner/v1/engine/train_engine.py index 74a56d4643..f1e9997604 100644 --- a/xtuner/v1/engine/train_engine.py +++ b/xtuner/v1/engine/train_engine.py @@ -30,6 +30,7 @@ from xtuner.v1.config import FSDPConfig, OptimConfig from xtuner.v1.data_proto.sequence_context import SequenceContext from xtuner.v1.loss import LogProbContext +from xtuner.v1.model.adapter.lora import LoraConfig from xtuner.v1.model.base import ( BaseModel, BatchForwardInfo, @@ -147,10 +148,12 @@ def __init__( optim_cfg: OptimConfig, fsdp_cfg: FSDPConfig, intra_layer_micro_batch: int = 1, + adapter_cfg: LoraConfig | None = None, ) -> None: self.model_cfg = model_cfg self.optim_cfg = optim_cfg self.fsdp_cfg = fsdp_cfg + self.adapter_cfg = adapter_cfg self.model = self.build_model() self.optimizer = self.build_optimizer(optim_cfg) self.intra_layer_micro_batch = intra_layer_micro_batch @@ -170,6 +173,8 @@ def __has_freeze_params(self) -> bool: def build_model(self) -> BaseModel: with torch.device("meta"): model = self.model_cfg.build() + if self.adapter_cfg: + model = self.adapter_cfg.build(model) model = model.fully_shard(self.fsdp_cfg) diff --git a/xtuner/v1/model/adapter/__init__.py b/xtuner/v1/model/adapter/__init__.py new file mode 100644 index 0000000000..e69de29bb2 diff --git a/xtuner/v1/model/adapter/lora.py b/xtuner/v1/model/adapter/lora.py new file mode 100644 index 0000000000..41e799a972 --- /dev/null +++ b/xtuner/v1/model/adapter/lora.py @@ -0,0 +1,382 @@ +import functools +import json +import os +from pathlib import Path +from typing import List, Optional, Union + +import torch +import torch.distributed as dist +import torch.nn as nn +from pydantic import BaseModel as PydanticBaseModel +from torch.distributed.tensor import DTensor +from tqdm import tqdm + +from xtuner.v1.model.base import _save_file +from xtuner.v1.module.grouped_linear.moe_group_linear import GroupedLinear +from xtuner.v1.module.lora_linear.lora_grouped_linear import LoraGroupedLinear +from xtuner.v1.module.lora_linear.lora_linear import LoraLinear +from xtuner.v1.utils import get_device, get_torch_device_module, profile_time_and_memory +from xtuner.v1.utils.load_spec import LoadSpec + + +DEVICE_MODULE = get_torch_device_module() +DEVICE = get_device() + + +class LoraConfig(PydanticBaseModel): + # 与 peft LoraConfig 对齐的字段 + r: int = 8 + target_modules: Optional[Union[List[str], str]] = None + lora_alpha: int = 8 + lora_dropout: float = 0.0 + bias: str = "none" # "none" | "all" | "lora_only" + modules_to_save: Optional[List[str]] = None + init_lora_weights: bool = True + layers_to_transform: Optional[Union[List[int], int]] = None # 暂不使用 + layers_pattern: Optional[str] = None # 暂不使用 + base_model_name_or_path: Optional[str] = None + + def build(self, base_model): + return LoraModel(base_model, self) + + def save_hf(self, hf_path: str | Path): + """Save the configuration to a HuggingFace-compatible format. + + Args: + hf_path (str | Path): Path where the configuration should be saved. + """ + hf_config = { + "alpha_pattern": {}, + "auto_mapping": None, + "base_model_name_or_path": self.base_model_name_or_path, + "bias": self.bias, + "corda_config": None, + "eva_config": None, + "exclude_modules": None, + "fan_in_fan_out": False, + "inference_mode": True, + "init_lora_weights": self.init_lora_weights, + "layer_replication": None, + "layers_pattern": self.layers_pattern, + "layers_to_transform": self.layers_to_transform, + "loftq_config": {}, + "lora_alpha": self.lora_alpha, + "lora_bias": False, + "lora_dropout": self.lora_dropout, + "megatron_config": None, + "megatron_core": "megatron.core", + "modules_to_save": self.modules_to_save, + "peft_type": "LORA", + "qalora_group_size": 16, + "r": self.r, + "rank_pattern": {}, + "revision": None, + "target_modules": self.target_modules, + "target_parameters": None, + "task_type": "CAUSAL_LM", + "trainable_token_indices": None, + "use_dora": False, + "use_qalora": False, + "use_rslora": False, + } + with open(os.path.join(hf_path, "adapter_config.json"), "w") as f: + json.dump(hf_config, f, indent=2) + + +def wrap_to_hf_key_list(obj): + orig = getattr(obj, "to_hf_key_list") + + @functools.wraps(orig) + def new_to_hf_key_list(key: str): + if "base_layer." in key: + key = key.replace("base_layer.", "") + out = orig(key) + # if ".base_layer." in out[0]: + # out[0] = out[0].replace(".base_layer.", ".") + return out + + setattr(obj, "to_hf_key_list", new_to_hf_key_list) + + +class LoraModel(nn.Module): + def __init__(self, model: nn.Module, lora_config: LoraConfig): + super().__init__() + self.base_model = model + self.lora_config = lora_config + + # apply lora to base model + self._apply_lora() + + # 5. 修改 hf 的 key mapping + wrap_to_hf_key_list(self.base_model) + self.base_model._init_load_spec() + + def _apply_lora(self): + # 1. 冻结整个原模型 + for p in self.base_model.parameters(): + p.requires_grad = False + + # 2. 注入 LoRA + self._replace_linear_layers(self.base_model, prefix="") + + # 3. 按 config.bias 设置 bias 的 requires_grad + self._apply_bias_setting() + + # 4. 按 modules_to_save 让特定模块参数仍然可训练（例如 lm_head） + self._apply_modules_to_save() + + # def __getattr__(self, name): + # try: + # return super().__getattr__(name) + # except AttributeError: + # return getattr(self.model, name) + + def _match_target(self, module_name: str) -> bool: + """与 peft 类似的逻辑： + + - target_modules=None: 所有 nn.Linear 都加 LoRA + - target_modules=str: 名字中包含该 substring 的模块 + - target_modules=List[str]: 名字中包含任一 substring 的模块 + """ + target = self.lora_config.target_modules + if target is None: + return True + if isinstance(target, str): + return target in module_name + return any(t in module_name for t in target) + + def _replace_linear_layers(self, module: nn.Module, prefix: str): + for name, child in list(module.named_children()): + full_name = f"{prefix}.{name}" if prefix else name + + if isinstance(child, (LoraLinear, LoraGroupedLinear)): + # 避免重复 wrap + continue + + if self._match_target(full_name) and isinstance(child, nn.Linear): + lora_layer = LoraLinear( + base_layer=child, + rank=self.lora_config.r, + alpha=self.lora_config.lora_alpha, + lora_dropout=self.lora_config.lora_dropout, + init_lora_weights=self.lora_config.init_lora_weights, + ) + setattr(module, name, lora_layer) + elif self._match_target(full_name) and isinstance(child, GroupedLinear): + lora_layer = LoraGroupedLinear( + base_layer=child, + rank=self.lora_config.r, + alpha=self.lora_config.lora_alpha, + lora_dropout=self.lora_config.lora_dropout, + init_lora_weights=self.lora_config.init_lora_weights, + ) + setattr(module, name, lora_layer) + else: + self._replace_linear_layers(child, prefix=full_name) + + def _apply_bias_setting(self): + """ + - "none": 所有 bias 冻结（默认） + - "all": 所有 bias 可训练 + - "lora_only": 只有挂了 LoRA 的 Linear 的 bias 可训练 + """ + bias_mode = self.lora_config.bias + + if bias_mode == "none": + # 已经在 init 时全部冻住，无需额外处理 + return + + if bias_mode == "all": + for module in self.base_model.modules(): + if isinstance(module, (nn.Linear, GroupedLinear)) and module.bias is not None: + module.bias.requires_grad = True + return + + if bias_mode == "lora_only": + for module in self.base_model.modules(): + if isinstance(module, (LoraLinear, LoraGroupedLinear)) and module.base_layer.bias is not None: + module.base_layer.bias.requires_grad = True + return + + raise ValueError(f"Unknown bias mode: {bias_mode}, expected one of ['none', 'all', 'lora_only']") + + def _apply_modules_to_save(self): + """modules_to_save 里的模块，即使用了 LoRA 也会直接训练 一般用来保留 lm_head、classification + head 等。""" + modules_to_save = self.lora_config.modules_to_save + if not modules_to_save: + return + + # 根据模块名匹配 + for module_name, module in self.base_model.named_modules(): + if module_name in modules_to_save: + for p in module.parameters(): + p.requires_grad = True + + def forward(self, *args, **kwargs): + return self.base_model(*args, **kwargs) + + @torch.no_grad() + def merge_lora(self): + """Merge LoRA weights into base layers. + + 主要用于 RL 权重更新场景：推理引擎目前只支持接收全量参数（base weight + LoRA delta）， + 不支持直接加载 LoRA 参数。因此在 weight update 前需要临时 merge，update 后再 unmerge。 + """ + for module in tqdm(self.base_model.modules(), desc="[Merge LoRA]"): + if isinstance(module, LoraLinear): + module.merge_lora() + + @torch.no_grad() + def unmerge_lora(self): + """Unmerge LoRA weights from base layers. + + 与 merge_lora() 配对使用，在 RL 权重更新完成后恢复 LoRA 状态， + 保证训练侧模型可以继续正常训练。 + """ + for module in tqdm(self.base_model.modules(), desc="[Unmerge LoRA]"): + if isinstance(module, LoraLinear): + module.unmerge_lora() + + def trainable_parameters(self): + params = [(name, param) for name, param in self.named_parameters() if param.requires_grad] + return params + + def print_trainable_parameters(self): + trainable_params = 0 + all_params = 0 + for _, param in self.named_parameters(): + num = param.numel() + all_params += num + if param.requires_grad: + trainable_params += num + + print( + f"trainable params: {trainable_params} || " + f"all params: {all_params} || " + f"trainable: {100 * trainable_params / all_params:.4f}%" + ) + + def fully_shard(self, *args, **kwargs): + self.base_model = self.base_model.fully_shard(*args, **kwargs) + return self + + @property + def device(self) -> torch.device: + return self.base_model.device + + def to_device(self, device: torch.device | str): + self.base_model.to_device(device) + + def scale_and_reduce_grad(self): + self.base_model.scale_and_reduce_grad() + + def set_hf(self, hf_path: str | Path): + self.lora_config.base_model_name_or_path = str(hf_path) + self.base_model.set_hf(hf_path) + + def save_hf(self, hf_dir: Path | str, save_dtype: torch.dtype = torch.bfloat16): + with profile_time_and_memory(f"[Saving HF to {hf_dir} cost]"): + self._save_hf(hf_dir=hf_dir, save_dtype=save_dtype) + + def _save_hf(self, hf_dir: Path | str, save_dtype: torch.dtype = torch.bfloat16): + """Save the hf model to the given directory. + + Args: + hf_dir (str): The directory to save the model. + save_dtype (torch.dtype): The dtype to save the model parameters, bfloat16 or float8. + """ + # if self._hf_path is None and self.config.hf_config is None: + # raise NotImplementedError( + # "The model is not loaded from Huggingface, and the `hf_config` property is not implemented, so it cannot be saved in Huggingface format." + # ) + + if isinstance(hf_dir, str): + hf_dir = Path(hf_dir) + hf_dir.mkdir(parents=True, exist_ok=True) + + DEVICE_MODULE.empty_cache() + assert save_dtype in [torch.float8_e4m3fn, torch.bfloat16], f"save_dtype {save_dtype} is not supported" + + lora_param = {} + modules_to_save = set(self.lora_config.modules_to_save or []) + + for name, param in self.state_dict().items(): + if "lora_A." in name or "lora_B." in name: + lora_param[name] = param + continue + + if self.lora_config.bias != "none" and name.endswith("bias"): + if ".base_layer.bias" in name: + lora_param[name] = param + continue + + if modules_to_save: + for mod in modules_to_save: + if mod in name: + lora_param[name] = param + + tensor_list: list[torch.Tensor] = [] + load_spec_list: list[LoadSpec] = [] + name_list: list[str] = [] + for name, param in lora_param.items(): + local_tensor = param._local_tensor if isinstance(param, DTensor) else param + local_tensor = local_tensor.bfloat16() + base_name = name[11:] # remove "base_model." prefix + if "lora_A." in name or "lora_B." in name: + load_spec = self.base_model.load_spec_mapping.get(base_name) + base_layer_name = base_name.replace("lora_A.", "base_layer.").replace("lora_B.", "base_layer.") + base_layer_load_spec = self.base_model.load_spec_mapping.get(base_layer_name) + hf_name = "base_model." + base_layer_load_spec.hf_keys[0][:-6] + "lora_" + name.split("lora_")[-1] + else: + load_spec = self.base_model.load_spec_mapping.get(base_name) + hf_name = "base_model." + load_spec.hf_keys[0] + tensor_list.append(local_tensor) + name_list.append(hf_name) + load_spec_list.append(load_spec) + + if self.base_model.fsdp_mesh is not None: + gathered_tensor_list = self.base_model._fsdp_foreach_allgather(tensor_list, load_spec_list) + else: + gathered_tensor_list = tensor_list + gathered_tensor_list = [ + self.base_model.param_to_safetensor(safetensor, name) + for safetensor, name in zip(gathered_tensor_list, name_list) + ] + + # Sepreately save fused parameters and others to make sure each saving rank will not save + # dupilicated keys + # + weight_map = {} + + safetensor_name = "adapter_model.safetensors" + + if not dist.is_initialized() or dist.get_rank() == 0: + # for tie_word_embeddings, we need to make sure each key is only saved once + unique_name_list = [] + unique_hf_tensor_list = [] + for name, hf_tensor in zip(name_list, gathered_tensor_list): + if name not in weight_map: + unique_name_list.append(name) + unique_hf_tensor_list.append(hf_tensor) + weight_map[name] = safetensor_name + + _save_file( + dict(zip(unique_name_list, unique_hf_tensor_list)), + hf_dir / safetensor_name, + ) + self.lora_config.save_hf(hf_dir) + + if dist.is_initialized(): + torch.distributed.barrier() + + def from_hf(self, hf_path: str | Path, strict: bool = True) -> tuple: + # TODO: load lora weight + loaded_keys, unloaded_keys, missing_keys = self.base_model.from_hf(hf_path, strict=False) + for unloaded_key in unloaded_keys: + assert "lora_A." in unloaded_key or "lora_B." in unloaded_key, ( + f"unloaded key {unloaded_key} is not a lora key" + ) + for missing_key in missing_keys: + assert "lora_A." in missing_key or "lora_B." in missing_key, f"missing key {missing_key} is not a lora key" + return loaded_keys, unloaded_keys, missing_keys \ No newline at end of file diff --git a/xtuner/v1/module/lora_linear/__init__.py b/xtuner/v1/module/lora_linear/__init__.py new file mode 100644 index 0000000000..e69de29bb2 diff --git a/xtuner/v1/module/lora_linear/lora_grouped_linear.py b/xtuner/v1/module/lora_linear/lora_grouped_linear.py new file mode 100644 index 0000000000..8827e35b33 --- /dev/null +++ b/xtuner/v1/module/lora_linear/lora_grouped_linear.py @@ -0,0 +1,97 @@ +import math + +import torch +import torch.nn as nn +from torch.distributed.tensor import DTensor + +from ..grouped_linear.moe_group_linear import GroupedLinear, build_grouped_linear + + +class LoraGroupedLinear(nn.Module): + def __init__( + self, + base_layer: GroupedLinear, + rank: int, + alpha: int, + lora_dropout: float = 0.0, + init_lora_weights: bool = True, + ): + super().__init__() + + self.base_layer = base_layer + self.in_features = base_layer.in_features + self.out_features = base_layer.out_features + self.num_routed_experts = base_layer.num_routed_experts + self.ep_mesh = base_layer.ep_mesh + self.rank = rank + self.alpha = alpha + self.scale = alpha / rank + self.merged = False + + self.lora_A = build_grouped_linear(self.in_features, self.rank, self.num_routed_experts, ep_mesh=self.ep_mesh) + self.lora_B = build_grouped_linear(self.rank, self.out_features, self.num_routed_experts, ep_mesh=self.ep_mesh) + + self.lora_dropout = nn.Dropout(p=lora_dropout) if lora_dropout > 0.0 else nn.Identity() + + for p in self.base_layer.parameters(): + p.requires_grad = False + + if init_lora_weights: + self.reset_parameters() + + def reset_parameters(self): + if isinstance(self.lora_A.weight, DTensor): + # TODO: init DTensor + raise NotImplementedError + else: + nn.init.kaiming_uniform_(self.lora_A.weight, a=math.sqrt(5)) + nn.init.zeros_(self.lora_B.weight) + + def forward(self, x: torch.Tensor, tokens_per_expert: torch.Tensor, decoding: bool = False) -> torch.Tensor: + if self.merged: + return self.base_layer(x, tokens_per_expert, decoding) + original_out = self.base_layer(x, tokens_per_expert, decoding) + # lora_out = self.lora_a_naive_forward(x, tokens_per_expert, decoding) + lora_out = self.lora_A(x, tokens_per_expert, decoding) + # lora_out = self.lora_b_naive_forward(lora_out, tokens_per_expert, decoding) + lora_out = self.lora_B(lora_out, tokens_per_expert, decoding) + return original_out + lora_out * self.scale + + def lora_a_naive_forward( + self, x: torch.Tensor, tokens_per_expert: torch.Tensor, decoding: bool = False + ) -> torch.Tensor: + weight = self.lora_A.weight.view(-1, self.rank, self.in_features) + batch_sizes = tokens_per_expert.cpu().numpy() + + out = [] + start = 0 + for i, size in enumerate(batch_sizes): + rhs = weight[i, :, :].t() + out.append(x[start : start + size, :] @ rhs) + start += size + return torch.cat(out) + + def lora_b_naive_forward( + self, x: torch.Tensor, tokens_per_expert: torch.Tensor, decoding: bool = False + ) -> torch.Tensor: + weight = self.lora_B.weight.view(-1, self.out_features, self.rank) + batch_sizes = tokens_per_expert.cpu().numpy() + + out = [] + start = 0 + for i, size in enumerate(batch_sizes): + rhs = weight[i, :, :].t() + out.append(x[start : start + size, :] @ rhs) + start += size + return torch.cat(out) + + @torch.no_grad() + def merge_lora(self): + raise NotImplementedError + + @torch.no_grad() + def unmerge_lora(self): + raise NotImplementedError + + def __repr__(self) -> str: + return "lora." + super().__repr__() \ No newline at end of file diff --git a/xtuner/v1/module/lora_linear/lora_linear.py b/xtuner/v1/module/lora_linear/lora_linear.py new file mode 100644 index 0000000000..9f2580ae37 --- /dev/null +++ b/xtuner/v1/module/lora_linear/lora_linear.py @@ -0,0 +1,100 @@ +import math + +import torch +import torch.nn as nn + +from ..linear.linear import build_linear + + +class LoraLinear(nn.Module): + def __init__( + self, + base_layer: nn.Linear, + rank: int, + alpha: int, + lora_dropout: float = 0.0, + init_lora_weights: bool = True, + ): + super().__init__() + + self.base_layer = base_layer + self.in_features = base_layer.in_features + self.out_features = base_layer.out_features + self.rank = rank + self.alpha = alpha + self.scale = alpha / rank + self.merged = False + + weight = base_layer.weight + dtype = weight.dtype + device = weight.device + + # A: (in_features -> r) + self.lora_A = build_linear(self.in_features, rank, bias=False, device=device, dtype=dtype) + # B: (r -> out_features) + self.lora_B = build_linear(rank, self.out_features, bias=False, device=device, dtype=dtype) + + self.lora_dropout = nn.Dropout(p=lora_dropout) if lora_dropout > 0.0 else nn.Identity() + + for p in self.base_layer.parameters(): + p.requires_grad = False + + if init_lora_weights: + self.reset_parameters() + + def reset_parameters(self): + nn.init.kaiming_uniform_(self.lora_A.weight, a=math.sqrt(5)) + nn.init.zeros_(self.lora_B.weight) + + def forward(self, x: torch.Tensor) -> torch.Tensor: + if self.merged: + return self.base_layer(x) + + original_out = self.base_layer(x) + lora_out = self.lora_B(self.lora_A(self.lora_dropout(x))) * self.scale + return original_out + lora_out + + @torch.no_grad() + def merge_lora(self): + """把 LoRA 权重合并进 base_layer.weight.""" + if self.merged: + return + + # delta_W = B @ A shape: [out, in] + delta_w = torch.matmul(self.lora_B.weight, self.lora_A.weight) + self.base_layer.weight += delta_w * self.scale + + self.merged = True + # 合并后 GoRA 参数可以不再训练 + for p in self.lora_A.parameters(): + p.requires_grad = False + for p in self.lora_B.parameters(): + p.requires_grad = False + + @torch.no_grad() + def unmerge_lora(self): + """从 base_layer.weight 中还原 LoRA（如果之前 merge 过）""" + if not self.merged: + return + + delta_w = torch.matmul(self.lora_B.weight, self.lora_A.weight) + self.base_layer.weight -= delta_w * self.scale + + self.merged = False + for p in self.lora_A.parameters(): + p.requires_grad = True + for p in self.lora_B.parameters(): + p.requires_grad = True + + def extra_repr(self) -> str: + return ( + f"in_features={self.in_features}, " + f"out_features={self.out_features}, " + f"r={self.rank}, " + f"lora_alpha={self.alpha}, " + f"scale={self.scale}, " + f"merged={self.merged}" + ) + + def __repr__(self) -> str: + return "lora." + super().__repr__() \ No newline at end of file diff --git a/xtuner/v1/rl/trainer/worker.py b/xtuner/v1/rl/trainer/worker.py index 3f4f122f8b..7890575929 100644 --- a/xtuner/v1/rl/trainer/worker.py +++ b/xtuner/v1/rl/trainer/worker.py @@ -39,6 +39,7 @@ from xtuner.v1.loss import BaseLossContext, CELossConfig, LogProbConfig from xtuner.v1.loss.ce_loss import CELossContext, LMHeadLossContext from xtuner.v1.loss.mtp_loss import MTPLossConfig, MTPLossContext +from xtuner.v1.model.adapter.lora import LoraConfig, LoraModel from xtuner.v1.model.base import BaseModel as XtunerBaseModel from xtuner.v1.model.base import ModelItem, TransformerConfig from xtuner.v1.model.compose.base import BaseComposeConfig, BaseComposeModel @@ -323,6 +324,7 @@ def _build_engine(self, worker_cfg: WorkerConfig) -> TrainEngine: optim_cfg=worker_cfg.optim_cfg, fsdp_cfg=worker_cfg.fsdp_cfg, model_cfg=worker_cfg.model_cfg, + adapter_cfg=worker_cfg.adapter_cfg, ) if worker_cfg.load_from is not None: engine.from_hf(worker_cfg.load_from) diff --git a/xtuner/v1/rl/weight_update/update_weighter.py b/xtuner/v1/rl/weight_update/update_weighter.py index db4558845e..28134dea81 100644 --- a/xtuner/v1/rl/weight_update/update_weighter.py +++ b/xtuner/v1/rl/weight_update/update_weighter.py @@ -2,6 +2,7 @@ from typing import Any +from xtuner.v1.model.adapter.lora import LoraModel from xtuner.v1.rl.rollout.worker import RolloutConfig from .data import ( @@ -76,8 +77,24 @@ def update_weights(self): f"backend={self.rollout_info.backend!r}." ) assert self.weight_iterator is not None, "Weight iterator is not initialized." + + model = self._engine.model + lora_model = None + if isinstance(model, LoraModel): + lora_model = model + # TODO: 当前 weight update 流程只支持发送全量参数（base weight + merged LoRA delta）， + # 不支持直接发送 LoRA 参数（lora_A/lora_B）到推理引擎。 + # 因此这里需要先把 LoRA merge 进 base weight，更新完成后再 unmerge 恢复。 + # 后续优化方向：weight iterator 区分 base weight 和 LoRA 参数，推理引擎支持 + # 直接加载 LoRA 参数，避免 merge/unmerge 的开销和精度损失。 + lora_model.merge_lora() + self._transport.update(self.weight_iterator) + if lora_model is not None: + # 恢复 LoRA 状态，保持训练侧模型可继续训练 + lora_model.unmerge_lora() + def _set_transport(self) -> None: rollout_info = self.rollout_info assert rollout_info is not None, "bind_rollout_weight_update() must be called before setting transport." diff --git a/xtuner/v1/rl/weight_update/weight_iterator.py b/xtuner/v1/rl/weight_update/weight_iterator.py index 55f236ba11..b9a2064076 100644 --- a/xtuner/v1/rl/weight_update/weight_iterator.py +++ b/xtuner/v1/rl/weight_update/weight_iterator.py @@ -235,11 +235,20 @@ def iter_hf_batches(self, submodule=None, final_update=False): for name_list, fused_param_list in fused_gen: state_dict = {name: param.detach() for name, param in zip(name_list, fused_param_list)} + # 过滤掉 LoRA 参数：当前推理引擎只接收全量参数（base + merged LoRA）， + # 不需要单独的 lora_A/lora_B 参数。LoRA 参数在训练侧保留，通过 merge/unmerge + # 在推理时临时合并进 base weight。 + for key in list(state_dict.keys()): + if "lora_" in key: + del state_dict[key] yield WeightUpdateBatch(state_dict, train_enable_ep=train_enable_ep, finished=False) del state_dict, name_list, fused_param_list for name_list, param_list in chain(same_gen, shard_gen): state_dict = {name: param.detach() for name, param in zip(name_list, param_list)} + for key in list(state_dict.keys()): + if "lora_" in key: + del state_dict[key] yield WeightUpdateBatch(state_dict, train_enable_ep=train_enable_ep, finished=False) del state_dict, name_list, param_list diff --git a/xtuner/v1/train/trainer.py b/xtuner/v1/train/trainer.py index e61f46b224..f96327f6db 100644 --- a/xtuner/v1/train/trainer.py +++ b/xtuner/v1/train/trainer.py @@ -45,6 +45,7 @@ from xtuner.v1.engine import TrainEngine from xtuner.v1.engine.train_engine import TrainStepInfo from xtuner.v1.loss import CELossConfig +from xtuner.v1.model.adapter.lora import LoraConfig from xtuner.v1.model.base import ModelItem, XTunerBaseModelConfig from xtuner.v1.model.moe.moe import MoEConfig from xtuner.v1.patch import ( @@ -405,6 +406,7 @@ class TrainerConfig(BaseModel): lr_cfg: LRConfig loss_cfg: CELossConfig = CELossConfig() fsdp_cfg: FSDPConfig | None = None + adapter_cfg: LoraConfig | None = None global_batch_size: int | None work_dir: Path | str | None = None log_dir: Path | str | None = None @@ -525,6 +527,7 @@ def __init__( model_cfg: XTunerBaseModelConfig, optim_cfg: OptimConfig, fsdp_cfg: FSDPConfig | None = FSDPConfig(), + adapter_cfg: LoraConfig | None = None, dataset_cfg: DatasetConfigList | None = None, # TODO: Removed in version 1.1.0 dataloader_cfg: DataloaderConfig, loss_cfg: CELossConfig | None = CELossConfig(), @@ -720,6 +723,7 @@ def __init__( model_config=model_cfg, optim_config=optim_cfg, fsdp_config=fsdp_cfg, + adapter_config=adapter_cfg, load_checkpoint_path=self._load_checkpoint_cfg.checkpoint_path, strict=strict_load, intra_layer_micro_batch=intra_layer_micro_batch, @@ -774,6 +778,7 @@ def from_config(cls, config: TrainerConfig) -> Self: model_cfg=config.model_cfg, optim_cfg=config.optim_cfg, fsdp_cfg=config.fsdp_cfg, + adapter_cfg=config.adapter_cfg, dataset_cfg=config.dataset_cfg, dataloader_cfg=config.dataloader_cfg, loss_cfg=config.loss_cfg, @@ -1115,6 +1120,7 @@ def build_engine( model_config: XTunerBaseModelConfig, optim_config: OptimConfig, fsdp_config: FSDPConfig, + adapter_config: LoraConfig, load_checkpoint_path: str | Path | None, intra_layer_micro_batch: int = 1, strict: bool = True, @@ -1126,6 +1132,7 @@ def build_engine( model_config (TransformerConfig | BaseComposeConfig): Model configuration. optim_config (OptimConfig): Optimizer configuration. fsdp_config (FSDPConfig): FSDP configuration for distributed training. + adapter_config (LoraConfig): LoRA adapter configuration. resume_cfg (ResumeConfig | None): Resume configuration for continuing training. intra_layer_micro_batch (int): Intra-layer micro batch size for gradient accumulation. strict (bool): Whether to strictly load model weights. @@ -1137,6 +1144,7 @@ def build_engine( optim_cfg=optim_config, fsdp_cfg=fsdp_config, model_cfg=model_config, + adapter_config=adapter_config, intra_layer_micro_batch=intra_layer_micro_batch, ) if model_path is not None and (model_config.dcp_ignore_frozen_params or load_checkpoint_path is None): From c8b90a23370dd89dabf9c765fc295f979850b953 Mon Sep 17 00:00:00 2001 From: RangiLyu Date: Fri, 10 Jul 2026 16:42:24 +0800 Subject: [PATCH 2/2] [Fix] Stabilize LoRA training and weight updates --- tests/engine/test_dense_lora_train_engine.py | 82 +-- tests/module/test_lora_linear.py | 137 ++++ xtuner/v1/model/adapter/__init__.py | 4 + xtuner/v1/model/adapter/lora.py | 681 ++++++++++++++---- xtuner/v1/module/lora_linear/__init__.py | 5 + .../module/lora_linear/lora_grouped_linear.py | 113 ++- xtuner/v1/module/lora_linear/lora_linear.py | 63 +- xtuner/v1/rl/trainer/worker.py | 3 +- xtuner/v1/rl/weight_update/update_weighter.py | 11 +- xtuner/v1/rl/weight_update/weight_iterator.py | 45 +- xtuner/v1/train/trainer.py | 4 +- 11 files changed, 908 insertions(+), 240 deletions(-) create mode 100644 tests/module/test_lora_linear.py diff --git a/tests/engine/test_dense_lora_train_engine.py b/tests/engine/test_dense_lora_train_engine.py index 80d23b612c..12a498ba4b 100644 --- a/tests/engine/test_dense_lora_train_engine.py +++ b/tests/engine/test_dense_lora_train_engine.py @@ -1,24 +1,25 @@ import os -import tempfile import shutil +import tempfile import time + import parametrize import torch import torch.distributed as dist -from xtuner._testing import DeterministicDDPTestCase -from transformers import AutoTokenizer +from torch.optim.lr_scheduler import LambdaLR -from xtuner.v1.model.moe.moe import SequenceContext -from xtuner.v1.model.dense.qwen3 import Qwen3Dense8BConfig -from xtuner.v1.model.base import ModelItem -from xtuner.v1.loss.ce_loss import CELossConfig, CELossContextInputItem -from xtuner.v1.config import FSDPConfig, LRConfig, AdamWConfig +from transformers import AutoTokenizer +from xtuner._testing import DeterministicDDPTestCase +from xtuner.v1.config import AdamWConfig, FSDPConfig, LRConfig from xtuner.v1.engine.train_engine import TrainEngine -from torch.optim.lr_scheduler import LambdaLR +from xtuner.v1.loss.ce_loss import CELossConfig +from xtuner.v1.model.adapter.lora import LoraConfig +from xtuner.v1.model.base import ModelItem +from xtuner.v1.model.dense.qwen3 import Qwen3Dense8BConfig +from xtuner.v1.model.moe.moe import SequenceContext from xtuner.v1.utils import pad_to_max_length from xtuner.v1.utils.device import get_device from xtuner.v1.utils.test_utils import init_data_mesh -from xtuner.v1.model.adapter.lora import LoraConfig # Qwen3 8B @@ -92,28 +93,21 @@ def warmup_fn(x): labels = pad_to_max_length(labels, -100, max_length=8192) losses = [] - data_mesh = None + sp_mesh = None if sp_size > 1: data_mesh = init_data_mesh(str(DEVICE), sp_size) + sp_mesh = data_mesh["sp"] for _ in range(10): seq_ctx = SequenceContext.from_input_ids((input_ids,), device=DEVICE) labels = labels.to(DEVICE) seq_ctx.num_padding = pack_len - seq_ctx_list = [seq_ctx] - loss_ctx_input_list: list[CELossContextInputItem] = [ - CELossContextInputItem(shifted_labels=labels) - ] - LossContext = loss_cfg.loss_ctx_cls - batches_loss_kwargs = LossContext.build_batches_loss_kwargs( - loss_ctx_input_list, - loss_cfg, - ) - loss_kwargs = batches_loss_kwargs[0] - loss_ctx = LossContext(loss_cfg, loss_kwargs) - seq_ctx = seq_ctx_list[0] - engine_input = [ModelItem(seq_ctx=seq_ctx, loss_ctx=loss_ctx)] - loss_log, _ = engine.train_step(engine_input) + if sp_mesh is not None: + seq_ctx = seq_ctx.split(sequence_parallel_mesh=sp_mesh) + loss_ctx = loss_cfg.build(data={"shifted_labels": labels}, sp_mesh=sp_mesh) + loss_ctx = loss_cfg.loss_ctx_cls.build_batches([loss_ctx])[0] + engine_input = [ModelItem(seq_ctx=seq_ctx, loss_ctx={"lm": loss_ctx})] + loss_log = engine.train_step(engine_input)["logs_info"] grad_norm = engine.clip_grad_norm() engine.step_optimizer(grad_norm) lr_scheduler.step() @@ -128,7 +122,7 @@ def warmup_fn(x): torch.cuda.empty_cache() try: dist.destroy_process_group(pg) - except: + except Exception: pass @parametrize.parametrize( @@ -179,21 +173,23 @@ def test_save_and_load(self, device, tp_size, hsdp_sharding_size): dist.barrier() time.sleep(1) - # engine2 = TrainEngine( - # model_cfg=moe_cfg, - # optim_cfg=optim_cfg, - # fsdp_cfg=fsdp_cfg, - # ) - # engine2.from_hf(hf_path=temp_dir) - - # state_dict = engine.model.state_dict() - # state_dict2 = engine2.model.state_dict() - # for key, val in state_dict.items(): - # val2 = state_dict2[key] - # val = val.full_tensor().bfloat16() - # val2 = val2.full_tensor().bfloat16() - # self.assertTrue(torch.equal(val, val2[:val.shape[0]]), - # f"Mismatch in {key} between bf16 and fp8, {val} and {val2[:val.shape[0]]}") + engine2 = TrainEngine( + model_cfg=moe_cfg, + optim_cfg=optim_cfg, + fsdp_cfg=fsdp_cfg, + adapter_cfg=adapter_cfg, + ) + engine2.from_hf(hf_path=temp_dir, strict=True) + + state_dict = engine.model.state_dict() + state_dict2 = engine2.model.state_dict() + for key, val in state_dict.items(): + if "lora_A." not in key and "lora_B." not in key and "lm_head" not in key: + continue + val2 = state_dict2[key] + val = val.full_tensor().bfloat16() + val2 = val2.full_tensor().bfloat16() + self.assertTrue(torch.equal(val, val2), f"Mismatch in {key} after adapter round trip") if dist.get_rank() == 0: shutil.rmtree(temp_dir) @@ -201,7 +197,7 @@ def test_save_and_load(self, device, tp_size, hsdp_sharding_size): torch.cuda.empty_cache() try: dist.destroy_process_group(pg) - except: + except Exception: pass @property @@ -210,4 +206,4 @@ def world_size(self) -> int: @property def destroy_pg_upon_exit(self) -> bool: - return False \ No newline at end of file + return False diff --git a/tests/module/test_lora_linear.py b/tests/module/test_lora_linear.py new file mode 100644 index 0000000000..72afa0cb15 --- /dev/null +++ b/tests/module/test_lora_linear.py @@ -0,0 +1,137 @@ +import torch +import torch.nn as nn + +from xtuner.v1.model.adapter.lora import LoraConfig, LoraModel +from xtuner.v1.module.grouped_linear.moe_group_linear import GroupedLinear +from xtuner.v1.module.lora_linear.lora_grouped_linear import LoraGroupedLinear +from xtuner.v1.module.lora_linear.lora_linear import LoraLinear +from xtuner.v1.utils.load_spec import LoadEnum, LoadSpec + + +def test_lora_linear_merge_round_trip(): + base_layer = nn.Linear(4, 3, bias=False) + lora_layer = LoraLinear(base_layer, rank=2, alpha=4) + x = torch.randn(5, 4) + + torch.testing.assert_close(lora_layer(x), base_layer(x)) + + nn.init.normal_(lora_layer.lora_A.weight) + nn.init.normal_(lora_layer.lora_B.weight) + expected = lora_layer(x) + base_weight = base_layer.weight.detach().clone() + + lora_layer.merge_lora() + torch.testing.assert_close(lora_layer(x), expected) + lora_layer.unmerge_lora() + torch.testing.assert_close(base_layer.weight, base_weight) + + +def test_lora_linear_can_initialize_after_meta_materialization(): + with torch.device("meta"): + lora_layer = LoraLinear(nn.Linear(4, 3, bias=False), rank=2, alpha=4) + + lora_layer.to_empty(device="cpu") + lora_layer.reset_parameters() + assert torch.count_nonzero(lora_layer.lora_A.weight) > 0 + assert torch.count_nonzero(lora_layer.lora_B.weight) == 0 + + +def test_lora_grouped_linear_merge_round_trip(): + base_layer = GroupedLinear(in_features=4, out_features=3, num_routed_experts=2) + lora_layer = LoraGroupedLinear(base_layer, rank=2, alpha=4) + + nn.init.normal_(base_layer.weight) + nn.init.normal_(lora_layer.lora_A.weight) + nn.init.normal_(lora_layer.lora_B.weight) + base_weight = base_layer.weight.detach().clone() + a = lora_layer.lora_A.weight.view(2, 2, 4) + b = lora_layer.lora_B.weight.view(2, 3, 2) + expected = base_weight.view(2, 3, 4) + torch.bmm(b, a) * lora_layer.scale + + lora_layer.merge_lora() + torch.testing.assert_close(base_layer.weight.view(2, 3, 4), expected) + lora_layer.unmerge_lora() + torch.testing.assert_close(base_layer.weight, base_weight) + + +class _ToyModel(nn.Module): + fsdp_mesh = None + + def __init__(self): + super().__init__() + self.q_proj = nn.Linear(4, 4) + self.lm_head = nn.Linear(4, 4) + self.forward_marker = object() + + def to_hf_key_list(self, key: str) -> list[str]: + return [key] + + def _init_load_spec(self): + self.load_spec_mapping = {} + for name, param in self.state_dict().items(): + hf_key = self.to_hf_key_list(name)[0] + self.load_spec_mapping[name] = LoadSpec( + name=name, + hf_keys=[hf_key], + shape=tuple(param.shape), + load_enum=LoadEnum.SAME, + ) + + @staticmethod + def _clean_param_name(name: str) -> str: + return name + + @staticmethod + def _load_same_hf_param(param, load_spec, checkpoint_loader): + tensor = checkpoint_loader.load(load_spec.hf_keys[0]) + if tensor is None: + return load_spec.hf_keys.copy() + with torch.no_grad(): + param.copy_(tensor) + return [] + + @staticmethod + def _load_fused_hf_param(*_): + raise AssertionError("Toy model has no fused parameters") + + @staticmethod + def _load_shard_hf_param(*_): + raise AssertionError("Toy model has no sharded parameters") + + def set_hf(self, hf_path): + self.hf_path = hf_path + + def forwarded_method(self): + return self.forward_marker + + +def test_lora_model_forwards_base_model_api_and_preserves_modules_to_save(): + base_model = _ToyModel() + model = LoraModel( + base_model, + LoraConfig(target_modules=["q_proj", "lm_head"], modules_to_save=["lm_head"]), + ) + + assert isinstance(base_model.q_proj, LoraLinear) + assert isinstance(base_model.lm_head, nn.Linear) + assert all(param.requires_grad for param in base_model.lm_head.parameters()) + assert model.forwarded_method() is base_model.forward_marker + + +def test_adapter_save_load_round_trip(tmp_path): + config = LoraConfig(target_modules=["q_proj"], modules_to_save=["lm_head"], base_model_name_or_path="base") + model = LoraModel(_ToyModel(), config) + nn.init.normal_(model.base_model.q_proj.lora_A.weight) + nn.init.normal_(model.base_model.q_proj.lora_B.weight) + nn.init.normal_(model.base_model.lm_head.weight) + nn.init.normal_(model.base_model.lm_head.bias) + model._save_hf(tmp_path) + + restored = LoraModel(_ToyModel(), config.model_copy(deep=True)) + restored._load_adapter(tmp_path, strict=True) + + expected = {name: param for name, param in model.base_model.named_parameters() if param.requires_grad} + actual = {name: param for name, param in restored.base_model.named_parameters() if param.requires_grad} + assert expected.keys() == actual.keys() + for name in expected: + torch.testing.assert_close(actual[name].bfloat16(), expected[name].bfloat16(), rtol=0, atol=0) diff --git a/xtuner/v1/model/adapter/__init__.py b/xtuner/v1/model/adapter/__init__.py index e69de29bb2..f6f460a553 100644 --- a/xtuner/v1/model/adapter/__init__.py +++ b/xtuner/v1/model/adapter/__init__.py @@ -0,0 +1,4 @@ +from .lora import LoraConfig, LoraModel + + +__all__ = ["LoraConfig", "LoraModel"] diff --git a/xtuner/v1/model/adapter/lora.py b/xtuner/v1/model/adapter/lora.py index 41e799a972..d885b78554 100644 --- a/xtuner/v1/model/adapter/lora.py +++ b/xtuner/v1/model/adapter/lora.py @@ -1,40 +1,49 @@ import functools import json import os +import re +from concurrent.futures import Future +from itertools import chain from pathlib import Path -from typing import List, Optional, Union +from typing import Any, Literal import torch import torch.distributed as dist import torch.nn as nn from pydantic import BaseModel as PydanticBaseModel +from pydantic import ConfigDict, Field +from safetensors import safe_open +from torch.distributed.fsdp import FSDPModule from torch.distributed.tensor import DTensor -from tqdm import tqdm +from xtuner.v1.float8.float8_gmm_tile_wise import TileWiseFloat8GroupedLinear from xtuner.v1.model.base import _save_file +from xtuner.v1.model.compose.base import BaseComposeModel +from xtuner.v1.module import LMHead from xtuner.v1.module.grouped_linear.moe_group_linear import GroupedLinear from xtuner.v1.module.lora_linear.lora_grouped_linear import LoraGroupedLinear from xtuner.v1.module.lora_linear.lora_linear import LoraLinear -from xtuner.v1.utils import get_device, get_torch_device_module, profile_time_and_memory -from xtuner.v1.utils.load_spec import LoadSpec +from xtuner.v1.utils import get_torch_device_module, profile_time_and_memory +from xtuner.v1.utils.load_spec import LoadEnum, LoadSpec +from xtuner.v1.utils.loader import download_model_from_hub DEVICE_MODULE = get_torch_device_module() -DEVICE = get_device() class LoraConfig(PydanticBaseModel): - # 与 peft LoraConfig 对齐的字段 - r: int = 8 - target_modules: Optional[Union[List[str], str]] = None - lora_alpha: int = 8 - lora_dropout: float = 0.0 - bias: str = "none" # "none" | "all" | "lora_only" - modules_to_save: Optional[List[str]] = None + model_config = ConfigDict(extra="forbid") + + r: int = Field(default=8, gt=0) + target_modules: list[str] | str | None = None + lora_alpha: int = Field(default=8, gt=0) + lora_dropout: float = Field(default=0.0, ge=0.0, lt=1.0) + bias: Literal["none", "all", "lora_only"] = "none" + modules_to_save: list[str] | None = None init_lora_weights: bool = True - layers_to_transform: Optional[Union[List[int], int]] = None # 暂不使用 - layers_pattern: Optional[str] = None # 暂不使用 - base_model_name_or_path: Optional[str] = None + layers_to_transform: list[int] | int | None = None # 暂不使用 + layers_pattern: str | None = None # 暂不使用 + base_model_name_or_path: str | None = None def build(self, base_model): return LoraModel(base_model, self) @@ -45,7 +54,7 @@ def save_hf(self, hf_path: str | Path): Args: hf_path (str | Path): Path where the configuration should be saved. """ - hf_config = { + hf_config: dict[str, Any] = { "alpha_pattern": {}, "auto_mapping": None, "base_model_name_or_path": self.base_model_name_or_path, @@ -80,7 +89,33 @@ def save_hf(self, hf_path: str | Path): "use_rslora": False, } with open(os.path.join(hf_path, "adapter_config.json"), "w") as f: - json.dump(hf_config, f, indent=2) + json.dump(hf_config, f, indent=2, ensure_ascii=False) + + +class _AdapterCheckpointLoader: + def __init__(self, adapter_path: Path): + self._file = safe_open(str(adapter_path), framework="pt") + self.weight_map = dict.fromkeys(self._file.keys(), adapter_path.name) + + def is_key_exist(self, key: str) -> bool: + return key in self.weight_map + + def load(self, key: str) -> torch.Tensor | None: + if key not in self.weight_map: + return None + return self._file.get_tensor(key) + + +class _TensorCheckpointLoader: + def __init__(self, tensors: dict[str, torch.Tensor]): + self.tensors = tensors + self.weight_map = dict.fromkeys(tensors, "adapter_model.safetensors") + + def is_key_exist(self, key: str) -> bool: + return key in self.tensors + + def load(self, key: str) -> torch.Tensor | None: + return self.tensors.get(key) def wrap_to_hf_key_list(obj): @@ -99,15 +134,17 @@ def new_to_hf_key_list(key: str): class LoraModel(nn.Module): + _PEFT_PREFIX = "base_model.model." + def __init__(self, model: nn.Module, lora_config: LoraConfig): super().__init__() + if isinstance(model, BaseComposeModel): + raise NotImplementedError("LoRA adapters for compose models are not supported yet") self.base_model = model self.lora_config = lora_config + self._loaded_adapter_path: Path | None = None - # apply lora to base model self._apply_lora() - - # 5. 修改 hf 的 key mapping wrap_to_hf_key_list(self.base_model) self.base_model._init_load_spec() @@ -125,35 +162,82 @@ def _apply_lora(self): # 4. 按 modules_to_save 让特定模块参数仍然可训练（例如 lm_head） self._apply_modules_to_save() - # def __getattr__(self, name): - # try: - # return super().__getattr__(name) - # except AttributeError: - # return getattr(self.model, name) - - def _match_target(self, module_name: str) -> bool: - """与 peft 类似的逻辑： + if not any(param.requires_grad for param in self.base_model.parameters()): + raise ValueError( + f"No trainable parameters matched target_modules={self.lora_config.target_modules} " + f"or modules_to_save={self.lora_config.modules_to_save}" + ) - - target_modules=None: 所有 nn.Linear 都加 LoRA - - target_modules=str: 名字中包含该 substring 的模块 - - target_modules=List[str]: 名字中包含任一 substring 的模块 - """ + def __getattr__(self, name: str): + try: + return super().__getattr__(name) + except AttributeError: + base_model = super().__getattr__("base_model") + return getattr(base_model, name) + + @staticmethod + def _matches_module_name(module_name: str, candidates: list[str]) -> bool: + return any(module_name == candidate or module_name.endswith(f".{candidate}") for candidate in candidates) + + def _is_module_to_save(self, module_name: str) -> bool: + modules_to_save = self.lora_config.modules_to_save or [] + return self._matches_module_name(module_name, modules_to_save) + + @staticmethod + def _target_candidate_names(module_name: str) -> list[str]: + candidates = [module_name] + if module_name.endswith(".fused_w1w3"): + prefix = module_name.removesuffix("fused_w1w3") + candidates.extend([prefix + "gate_proj", prefix + "up_proj"]) + elif module_name.endswith(".fused_w2"): + candidates.append(module_name.removesuffix("fused_w2") + "down_proj") + return candidates + + def _raw_target_match(self, module_name: str) -> bool: target = self.lora_config.target_modules if target is None: return True if isinstance(target, str): - return target in module_name - return any(t in module_name for t in target) + if target == "all-linear": + return True + return re.fullmatch(target, module_name) is not None + return self._matches_module_name(module_name, target) + + def _match_target(self, module_name: str) -> bool: + """与 peft 类似的逻辑： + + - target_modules=None: 所有支持的 Linear（不包含 LMHead） + - target_modules=str: 按正则表达式完整匹配 + - target_modules=List[str]: 匹配完整模块名或模块名后缀 + """ + candidate_names = self._target_candidate_names(module_name) + target_matches = any(self._raw_target_match(candidate) for candidate in candidate_names) + + if not target_matches or self.lora_config.layers_to_transform is None: + return target_matches + + layer_indices = self.lora_config.layers_to_transform + if isinstance(layer_indices, int): + layer_indices = [layer_indices] + layer_pattern = self.lora_config.layers_pattern or "layers" + match = re.search(rf"(?:^|\.){re.escape(layer_pattern)}\.(\d+)(?:\.|$)", module_name) + return match is not None and int(match.group(1)) in layer_indices def _replace_linear_layers(self, module: nn.Module, prefix: str): for name, child in list(module.named_children()): full_name = f"{prefix}.{name}" if prefix else name if isinstance(child, (LoraLinear, LoraGroupedLinear)): - # 避免重复 wrap + continue + + if self._is_module_to_save(full_name): continue if self._match_target(full_name) and isinstance(child, nn.Linear): + if isinstance(child, LMHead): + if self.lora_config.target_modules is None: + continue + raise ValueError("LMHead cannot be wrapped by LoraLinear; add it to modules_to_save instead") lora_layer = LoraLinear( base_layer=child, rank=self.lora_config.r, @@ -162,7 +246,15 @@ def _replace_linear_layers(self, module: nn.Module, prefix: str): init_lora_weights=self.lora_config.init_lora_weights, ) setattr(module, name, lora_layer) - elif self._match_target(full_name) and isinstance(child, GroupedLinear): + elif self._match_target(full_name) and isinstance(child, (GroupedLinear, TileWiseFloat8GroupedLinear)): + if full_name.endswith(".fused_w1w3") and not self._raw_target_match(full_name): + prefix = full_name.removesuffix("fused_w1w3") + gate_matched = self._raw_target_match(prefix + "gate_proj") + up_matched = self._raw_target_match(prefix + "up_proj") + if gate_matched != up_matched: + raise ValueError( + "fused_w1w3 uses a shared LoRA A matrix; gate_proj and up_proj must be targeted together" + ) lora_layer = LoraGroupedLinear( base_layer=child, rank=self.lora_config.r, @@ -187,19 +279,19 @@ def _apply_bias_setting(self): return if bias_mode == "all": - for module in self.base_model.modules(): - if isinstance(module, (nn.Linear, GroupedLinear)) and module.bias is not None: - module.bias.requires_grad = True + for name, param in self.base_model.named_parameters(): + if name.endswith("bias"): + param.requires_grad = True return if bias_mode == "lora_only": for module in self.base_model.modules(): - if isinstance(module, (LoraLinear, LoraGroupedLinear)) and module.base_layer.bias is not None: - module.base_layer.bias.requires_grad = True + if isinstance(module, (LoraLinear, LoraGroupedLinear)): + bias = getattr(module.base_layer, "bias", None) + if bias is not None: + bias.requires_grad = True return - raise ValueError(f"Unknown bias mode: {bias_mode}, expected one of ['none', 'all', 'lora_only']") - def _apply_modules_to_save(self): """modules_to_save 里的模块，即使用了 LoRA 也会直接训练 一般用来保留 lm_head、classification head 等。""" @@ -207,15 +299,65 @@ def _apply_modules_to_save(self): if not modules_to_save: return - # 根据模块名匹配 for module_name, module in self.base_model.named_modules(): - if module_name in modules_to_save: + if self._matches_module_name(module_name, modules_to_save): for p in module.parameters(): p.requires_grad = True def forward(self, *args, **kwargs): return self.base_model(*args, **kwargs) + def init_weights(self): + self.base_model.init_weights() + self._reset_lora_parameters() + + def _reset_lora_parameters(self): + for module in self.base_model.modules(): + if isinstance(module, (LoraLinear, LoraGroupedLinear)): + module.reset_parameters() + + def _lora_modules_by_fsdp_manager( + self, + ) -> tuple[list[tuple[FSDPModule, list[LoraLinear | LoraGroupedLinear]]], list[LoraLinear | LoraGroupedLinear]]: + named_modules = dict(self.base_model.named_modules()) + managed: dict[int, tuple[FSDPModule, list[LoraLinear | LoraGroupedLinear]]] = {} + standalone: list[LoraLinear | LoraGroupedLinear] = [] + + for module_name, module in named_modules.items(): + if not isinstance(module, (LoraLinear, LoraGroupedLinear)): + continue + + manager = None + prefix = module_name + while True: + candidate = named_modules[prefix] + if isinstance(candidate, FSDPModule): + manager = candidate + break + if not prefix: + break + prefix = prefix.rpartition(".")[0] + + if manager is None: + standalone.append(module) + else: + managed.setdefault(id(manager), (manager, []))[1].append(module) + + return list(managed.values()), standalone + + def _apply_lora_weight_op(self, op_name: Literal["merge_lora", "unmerge_lora"]): + managed, standalone = self._lora_modules_by_fsdp_manager() + for manager, modules in managed: + manager.unshard() + try: + for module in modules: + getattr(module, op_name)() + finally: + manager.reshard() + + for module in standalone: + getattr(module, op_name)() + @torch.no_grad() def merge_lora(self): """Merge LoRA weights into base layers. @@ -223,9 +365,16 @@ def merge_lora(self): 主要用于 RL 权重更新场景：推理引擎目前只支持接收全量参数（base weight + LoRA delta）， 不支持直接加载 LoRA 参数。因此在 weight update 前需要临时 merge，update 后再 unmerge。 """ - for module in tqdm(self.base_model.modules(), desc="[Merge LoRA]"): - if isinstance(module, LoraLinear): - module.merge_lora() + try: + self._apply_lora_weight_op("merge_lora") + except BaseException: + if any( + module.merged + for module in self.base_model.modules() + if isinstance(module, (LoraLinear, LoraGroupedLinear)) + ): + self._apply_lora_weight_op("unmerge_lora") + raise @torch.no_grad() def unmerge_lora(self): @@ -234,9 +383,7 @@ def unmerge_lora(self): 与 merge_lora() 配对使用，在 RL 权重更新完成后恢复 LoRA 状态， 保证训练侧模型可以继续正常训练。 """ - for module in tqdm(self.base_model.modules(), desc="[Unmerge LoRA]"): - if isinstance(module, LoraLinear): - module.unmerge_lora() + self._apply_lora_weight_op("unmerge_lora") def trainable_parameters(self): params = [(name, param) for name, param in self.named_parameters() if param.requires_grad] @@ -259,6 +406,11 @@ def print_trainable_parameters(self): def fully_shard(self, *args, **kwargs): self.base_model = self.base_model.fully_shard(*args, **kwargs) + if self.base_model.fsdp_config is None or self.base_model.fsdp_config.requires_grad: + # Some model implementations re-tie or replace parameters while sharding. + # Restore the adapter trainability policy before the optimizer is built. + self._apply_bias_setting() + self._apply_modules_to_save() return self @property @@ -272,111 +424,366 @@ def scale_and_reduce_grad(self): self.base_model.scale_and_reduce_grad() def set_hf(self, hf_path: str | Path): - self.lora_config.base_model_name_or_path = str(hf_path) - self.base_model.set_hf(hf_path) + hf_path_obj = Path(hf_path) if isinstance(hf_path, (str, Path)) else None + if hf_path_obj is not None and hf_path_obj.is_dir() and (hf_path_obj / "adapter_config.json").is_file(): + with open(hf_path_obj / "adapter_config.json") as f: + adapter_config = json.load(f) + base_model_path = adapter_config.get("base_model_name_or_path") + if base_model_path is None: + raise ValueError(f"Missing base_model_name_or_path in {hf_path_obj / 'adapter_config.json'}") + else: + base_model_path = str(hf_path) + + self.lora_config.base_model_name_or_path = str(base_model_path) + self.base_model.set_hf(base_model_path) + + def _get_load_spec(self, param_name: str) -> LoadSpec: + load_spec = self.base_model.load_spec_mapping.get(param_name) + if load_spec is None: + raise RuntimeError(f"Parameter {param_name} is missing from the base model load_spec_mapping") + return load_spec + + @staticmethod + def _all_gather_tensor(tensor: torch.Tensor, group: dist.ProcessGroup, dim: int) -> torch.Tensor: + gathered = [torch.empty_like(tensor) for _ in range(dist.get_world_size(group=group))] + dist.all_gather(gathered, tensor.contiguous(), group=group) + return torch.cat(gathered, dim=dim) + + @staticmethod + def _global_hf_keys(load_spec: LoadSpec) -> list[str]: + if load_spec.load_enum != LoadEnum.FUSED or load_spec.group is None or not dist.is_initialized(): + return load_spec.hf_keys.copy() + + gathered_keys: list[list[str] | None] = [None] * dist.get_world_size(group=load_spec.group) + dist.all_gather_object(gathered_keys, load_spec.hf_keys, group=load_spec.group) + return list(chain.from_iterable(keys for keys in gathered_keys if keys is not None)) + + def _gather_param(self, param: torch.Tensor, load_spec: LoadSpec) -> torch.Tensor: + local_tensor = param.to_local() if isinstance(param, DTensor) else param + + if self.base_model.fsdp_mesh is not None and tuple(local_tensor.shape) != tuple(load_spec.shape): + gathered_tensor = self.base_model._fsdp_foreach_allgather([local_tensor], [load_spec])[0] + else: + gathered_tensor = local_tensor + + if ( + load_spec.group is not None + and dist.is_initialized() + and load_spec.load_enum in (LoadEnum.FUSED, LoadEnum.SHARD) + ): + gathered_tensor = self._all_gather_tensor(gathered_tensor, load_spec.group, load_spec.dim or 0) + + return gathered_tensor.detach().to(dtype=torch.bfloat16, device="cpu").contiguous() + + @staticmethod + def _split_fused_tensor(tensor: torch.Tensor, load_spec: LoadSpec, hf_keys: list[str]) -> list[torch.Tensor]: + if len(hf_keys) == 1: + return [tensor] + dim = load_spec.dim or 0 + if tensor.shape[dim] % len(hf_keys) != 0: + raise RuntimeError( + f"Cannot split {load_spec.name} with shape {tuple(tensor.shape)} into {len(hf_keys)} HF tensors" + ) + return list(torch.chunk(tensor, len(hf_keys), dim=dim)) + + @classmethod + def _adapter_key(cls, hf_weight_key: str, lora_name: Literal["lora_A", "lora_B"] | None = None) -> str: + if lora_name is None: + return cls._PEFT_PREFIX + hf_weight_key + if not hf_weight_key.endswith("weight"): + raise RuntimeError(f"LoRA can only be exported for weight tensors, got {hf_weight_key}") + return cls._PEFT_PREFIX + hf_weight_key[:-6] + f"{lora_name}.weight" + + def _export_lora_module( + self, + module_name: str, + module: LoraLinear | LoraGroupedLinear, + adapter_state: dict[str, torch.Tensor], + ) -> None: + module_name = self.base_model._clean_param_name(module_name) + a_spec = self._get_load_spec(f"{module_name}.lora_A.weight") + b_spec = self._get_load_spec(f"{module_name}.lora_B.weight") + base_spec = self._get_load_spec(f"{module_name}.base_layer.weight") + base_hf_keys = self._global_hf_keys(base_spec) + lora_a = self._gather_param(module.lora_A.weight, a_spec) + lora_b = self._gather_param(module.lora_B.weight, b_spec) + + if isinstance(module, LoraLinear): + if len(base_hf_keys) != 1: + raise RuntimeError(f"Expected one HF key for {module_name}, got {base_hf_keys}") + adapter_state[self._adapter_key(base_hf_keys[0], "lora_A")] = lora_a + adapter_state[self._adapter_key(base_hf_keys[0], "lora_B")] = lora_b + return + + num_experts = module.num_routed_experts + if len(base_hf_keys) % num_experts != 0: + raise RuntimeError(f"HF keys for {module_name} cannot be divided across {num_experts} experts") + projections_per_expert = len(base_hf_keys) // num_experts + if module.out_features % projections_per_expert != 0: + raise RuntimeError(f"out_features of {module_name} cannot be split across its HF projection keys") + + a_view = lora_a.view(num_experts, module.rank, module.in_features) + b_view = lora_b.view(num_experts, module.out_features, module.rank) + b_chunks = b_view.chunk(projections_per_expert, dim=1) + for expert_idx in range(num_experts): + for projection_idx in range(projections_per_expert): + hf_key = base_hf_keys[expert_idx * projections_per_expert + projection_idx] + adapter_state[self._adapter_key(hf_key, "lora_A")] = a_view[expert_idx].contiguous() + adapter_state[self._adapter_key(hf_key, "lora_B")] = b_chunks[projection_idx][expert_idx].contiguous() + + def _adapter_state_dict(self) -> dict[str, torch.Tensor]: + adapter_state: dict[str, torch.Tensor] = {} + for module_name, module in self.base_model.named_modules(): + if isinstance(module, (LoraLinear, LoraGroupedLinear)): + self._export_lora_module(module_name, module, adapter_state) + + for param_name, param in self.base_model.named_parameters(): + param_name = self.base_model._clean_param_name(param_name) + if "lora_A." in param_name or "lora_B." in param_name or not param.requires_grad: + continue + + load_spec = self._get_load_spec(param_name) + hf_keys = self._global_hf_keys(load_spec) + gathered_param = self._gather_param(param, load_spec) + for hf_key, tensor in zip(hf_keys, self._split_fused_tensor(gathered_param, load_spec, hf_keys)): + adapter_state[self._adapter_key(hf_key)] = tensor.contiguous() + + return adapter_state def save_hf(self, hf_dir: Path | str, save_dtype: torch.dtype = torch.bfloat16): with profile_time_and_memory(f"[Saving HF to {hf_dir} cost]"): self._save_hf(hf_dir=hf_dir, save_dtype=save_dtype) def _save_hf(self, hf_dir: Path | str, save_dtype: torch.dtype = torch.bfloat16): - """Save the hf model to the given directory. - - Args: - hf_dir (str): The directory to save the model. - save_dtype (torch.dtype): The dtype to save the model parameters, bfloat16 or float8. - """ - # if self._hf_path is None and self.config.hf_config is None: - # raise NotImplementedError( - # "The model is not loaded from Huggingface, and the `hf_config` property is not implemented, so it cannot be saved in Huggingface format." - # ) - + """Save a PEFT-compatible adapter checkpoint.""" if isinstance(hf_dir, str): hf_dir = Path(hf_dir) hf_dir.mkdir(parents=True, exist_ok=True) - DEVICE_MODULE.empty_cache() + if hasattr(DEVICE_MODULE, "empty_cache"): + DEVICE_MODULE.empty_cache() assert save_dtype in [torch.float8_e4m3fn, torch.bfloat16], f"save_dtype {save_dtype} is not supported" + adapter_state = self._adapter_state_dict() - lora_param = {} - modules_to_save = set(self.lora_config.modules_to_save or []) + if not dist.is_initialized() or dist.get_rank() == 0: + _save_file(adapter_state, hf_dir / "adapter_model.safetensors") + self.lora_config.save_hf(hf_dir) - for name, param in self.state_dict().items(): - if "lora_A." in name or "lora_B." in name: - lora_param[name] = param - continue + if dist.is_initialized(): + torch.distributed.barrier() - if self.lora_config.bias != "none" and name.endswith("bias"): - if ".base_layer.bias" in name: - lora_param[name] = param - continue - - if modules_to_save: - for mod in modules_to_save: - if mod in name: - lora_param[name] = param - - tensor_list: list[torch.Tensor] = [] - load_spec_list: list[LoadSpec] = [] - name_list: list[str] = [] - for name, param in lora_param.items(): - local_tensor = param._local_tensor if isinstance(param, DTensor) else param - local_tensor = local_tensor.bfloat16() - base_name = name[11:] # remove "base_model." prefix - if "lora_A." in name or "lora_B." in name: - load_spec = self.base_model.load_spec_mapping.get(base_name) - base_layer_name = base_name.replace("lora_A.", "base_layer.").replace("lora_B.", "base_layer.") - base_layer_load_spec = self.base_model.load_spec_mapping.get(base_layer_name) - hf_name = "base_model." + base_layer_load_spec.hf_keys[0][:-6] + "lora_" + name.split("lora_")[-1] - else: - load_spec = self.base_model.load_spec_mapping.get(base_name) - hf_name = "base_model." + load_spec.hf_keys[0] - tensor_list.append(local_tensor) - name_list.append(hf_name) - load_spec_list.append(load_spec) - - if self.base_model.fsdp_mesh is not None: - gathered_tensor_list = self.base_model._fsdp_foreach_allgather(tensor_list, load_spec_list) - else: - gathered_tensor_list = tensor_list - gathered_tensor_list = [ - self.base_model.param_to_safetensor(safetensor, name) - for safetensor, name in zip(gathered_tensor_list, name_list) - ] + def async_save_hf(self, hf_dir: Path | str, save_dtype: torch.dtype = torch.bfloat16, **_: Any) -> Future[Path]: + future: Future[Path] = Future() + try: + self.save_hf(hf_dir=hf_dir, save_dtype=save_dtype) + future.set_result(Path(hf_dir)) + except BaseException as exc: + future.set_exception(exc) + return future + + def destroy_async_hf_resources(self) -> None: + return + + @staticmethod + def _resolve_adapter_dir(hf_path: str | Path) -> Path | None: + resolved_path = Path(download_model_from_hub(str(hf_path))) + if (resolved_path / "adapter_config.json").is_file() and ( + resolved_path / "adapter_model.safetensors" + ).is_file(): + return resolved_path + return None + + def _load_base_hf(self, hf_path: str | Path, strict: bool) -> tuple[set[str], set[str], set[str]]: + loaded_keys, unloaded_keys, missing_keys = self.base_model.from_hf(hf_path, strict=False) + self._reset_lora_parameters() - # Sepreately save fused parameters and others to make sure each saving rank will not save - # dupilicated keys - # - weight_map = {} + unexpected_unloaded = {key for key in unloaded_keys if "lora_A." not in key and "lora_B." not in key} + unexpected_missing = {key for key in missing_keys if "lora_A." not in key and "lora_B." not in key} + if strict and (unexpected_unloaded or unexpected_missing): + raise RuntimeError( + f"Failed to load base model. unloaded={unexpected_unloaded}, missing={unexpected_missing}" + ) - safetensor_name = "adapter_model.safetensors" + self.lora_config.base_model_name_or_path = str(hf_path) + return loaded_keys, unloaded_keys, missing_keys + + @classmethod + def _load_adapter_tensor( + cls, loader: _AdapterCheckpointLoader, hf_key: str, consumed: set[str] + ) -> torch.Tensor | None: + standard_key = cls._PEFT_PREFIX + hf_key + legacy_key = "base_model." + hf_key + for key in (standard_key, legacy_key): + tensor = loader.load(key) + if tensor is not None: + consumed.add(key) + return tensor + return None + + def _load_lora_module_tensors( + self, + module_name: str, + module: LoraLinear | LoraGroupedLinear, + loader: _AdapterCheckpointLoader, + consumed: set[str], + missing: set[str], + ) -> dict[str, torch.Tensor]: + module_name = self.base_model._clean_param_name(module_name) + a_spec = self._get_load_spec(f"{module_name}.lora_A.weight") + b_spec = self._get_load_spec(f"{module_name}.lora_B.weight") + base_spec = self._get_load_spec(f"{module_name}.base_layer.weight") + base_hf_keys = self._global_hf_keys(base_spec) + + if isinstance(module, LoraLinear): + if len(base_hf_keys) != 1: + raise RuntimeError(f"Expected one HF key for {module_name}, got {base_hf_keys}") + a_key = base_hf_keys[0][:-6] + "lora_A.weight" + b_key = base_hf_keys[0][:-6] + "lora_B.weight" + lora_a = self._load_adapter_tensor(loader, a_key, consumed) + lora_b = self._load_adapter_tensor(loader, b_key, consumed) + if lora_a is None: + missing.add(self._PEFT_PREFIX + a_key) + if lora_b is None: + missing.add(self._PEFT_PREFIX + b_key) + if lora_a is None or lora_b is None: + return {} + return {a_spec.hf_keys[0]: lora_a, b_spec.hf_keys[0]: lora_b} + + num_experts = module.num_routed_experts + projections_per_expert = len(base_hf_keys) // num_experts + expert_a: list[torch.Tensor] = [] + expert_b: list[torch.Tensor] = [] + for expert_idx in range(num_experts): + projection_a: list[torch.Tensor] = [] + projection_b: list[torch.Tensor] = [] + for projection_idx in range(projections_per_expert): + hf_key = base_hf_keys[expert_idx * projections_per_expert + projection_idx] + a_key = hf_key[:-6] + "lora_A.weight" + b_key = hf_key[:-6] + "lora_B.weight" + lora_a = self._load_adapter_tensor(loader, a_key, consumed) + lora_b = self._load_adapter_tensor(loader, b_key, consumed) + if lora_a is None: + missing.add(self._PEFT_PREFIX + a_key) + else: + projection_a.append(lora_a) + if lora_b is None: + missing.add(self._PEFT_PREFIX + b_key) + else: + projection_b.append(lora_b) + + if len(projection_a) != projections_per_expert or len(projection_b) != projections_per_expert: + continue + if any(not torch.equal(projection_a[0], tensor) for tensor in projection_a[1:]): + raise RuntimeError( + f"{module_name} shares lora_A across fused projections, but the adapter contains different A tensors" + ) + expert_a.append(projection_a[0]) + expert_b.append(torch.cat(projection_b, dim=0)) + + if len(expert_a) != num_experts or len(expert_b) != num_experts: + return {} + lora_a = torch.stack(expert_a).reshape(num_experts * module.rank, module.in_features) + lora_b = torch.stack(expert_b).reshape(num_experts * module.out_features, module.rank) + return {a_spec.hf_keys[0]: lora_a, b_spec.hf_keys[0]: lora_b} + + def _load_param_from_tensors( + self, param: torch.Tensor, load_spec: LoadSpec, tensor_loader: _TensorCheckpointLoader + ) -> list[str]: + if load_spec.load_enum == LoadEnum.SAME: + return self.base_model._load_same_hf_param(param, load_spec, tensor_loader) + if load_spec.load_enum == LoadEnum.FUSED: + return self.base_model._load_fused_hf_param(param, load_spec, tensor_loader) + if load_spec.load_enum == LoadEnum.SHARD: + return self.base_model._load_shard_hf_param(param, load_spec, tensor_loader) + raise RuntimeError(f"Unsupported load enum {load_spec.load_enum}") + + def _load_adapter(self, adapter_dir: Path, strict: bool) -> tuple[set[str], set[str], set[str]]: + loader = _AdapterCheckpointLoader(adapter_dir / "adapter_model.safetensors") + consumed: set[str] = set() + missing_adapter_keys: set[str] = set() + internal_tensors: dict[str, torch.Tensor] = {} + targets: list[tuple[str, torch.Tensor, LoadSpec]] = [] - if not dist.is_initialized() or dist.get_rank() == 0: - # for tie_word_embeddings, we need to make sure each key is only saved once - unique_name_list = [] - unique_hf_tensor_list = [] - for name, hf_tensor in zip(name_list, gathered_tensor_list): - if name not in weight_map: - unique_name_list.append(name) - unique_hf_tensor_list.append(hf_tensor) - weight_map[name] = safetensor_name - - _save_file( - dict(zip(unique_name_list, unique_hf_tensor_list)), - hf_dir / safetensor_name, + for module_name, module in self.base_model.named_modules(): + if not isinstance(module, (LoraLinear, LoraGroupedLinear)): + continue + internal_tensors.update( + self._load_lora_module_tensors( + module_name, module, loader, consumed=consumed, missing=missing_adapter_keys + ) + ) + clean_name = self.base_model._clean_param_name(module_name) + targets.extend( + [ + ( + f"{clean_name}.lora_A.weight", + module.lora_A.weight, + self._get_load_spec(f"{clean_name}.lora_A.weight"), + ), + ( + f"{clean_name}.lora_B.weight", + module.lora_B.weight, + self._get_load_spec(f"{clean_name}.lora_B.weight"), + ), + ] ) - self.lora_config.save_hf(hf_dir) - if dist.is_initialized(): - torch.distributed.barrier() + for param_name, param in self.base_model.named_parameters(): + param_name = self.base_model._clean_param_name(param_name) + if "lora_A." in param_name or "lora_B." in param_name or not param.requires_grad: + continue + load_spec = self._get_load_spec(param_name) + for hf_key in self._global_hf_keys(load_spec): + tensor = self._load_adapter_tensor(loader, hf_key, consumed) + if tensor is None: + missing_adapter_keys.add(self._adapter_key(hf_key)) + else: + internal_tensors[hf_key] = tensor + targets.append((param_name, param, load_spec)) + + tensor_loader = _TensorCheckpointLoader(internal_tensors) + loaded_params: set[str] = set() + unloaded_params: set[str] = set() + for param_name, param, load_spec in targets: + missing_internal = self._load_param_from_tensors(param, load_spec, tensor_loader) + if missing_internal: + unloaded_params.add(param_name) + else: + loaded_params.add(param_name) + + unexpected_keys = set(loader.weight_map) - consumed + if strict and (missing_adapter_keys or unloaded_params or unexpected_keys): + raise RuntimeError( + "Failed to load LoRA adapter. " + f"missing={missing_adapter_keys}, unloaded={unloaded_params}, unexpected={unexpected_keys}" + ) + + self._loaded_adapter_path = adapter_dir + return loaded_params, unloaded_params, missing_adapter_keys def from_hf(self, hf_path: str | Path, strict: bool = True) -> tuple: - # TODO: load lora weight - loaded_keys, unloaded_keys, missing_keys = self.base_model.from_hf(hf_path, strict=False) - for unloaded_key in unloaded_keys: - assert "lora_A." in unloaded_key or "lora_B." in unloaded_key, ( - f"unloaded key {unloaded_key} is not a lora key" + adapter_dir = self._resolve_adapter_dir(hf_path) + if adapter_dir is None: + return self._load_base_hf(hf_path, strict=strict) + + with open(adapter_dir / "adapter_config.json") as f: + adapter_config = json.load(f) + if adapter_config.get("r") != self.lora_config.r: + raise ValueError( + f"Adapter rank {adapter_config.get('r')} does not match configured rank {self.lora_config.r}" + ) + if adapter_config.get("lora_alpha") != self.lora_config.lora_alpha: + raise ValueError( + "Adapter lora_alpha " + f"{adapter_config.get('lora_alpha')} does not match configured value {self.lora_config.lora_alpha}" ) - for missing_key in missing_keys: - assert "lora_A." in missing_key or "lora_B." in missing_key, f"missing key {missing_key} is not a lora key" - return loaded_keys, unloaded_keys, missing_keys \ No newline at end of file + base_model_path = adapter_config.get("base_model_name_or_path") + if not base_model_path: + raise ValueError(f"Missing base_model_name_or_path in {adapter_dir / 'adapter_config.json'}") + + relative_base_path = adapter_dir / base_model_path + if not Path(base_model_path).is_absolute() and relative_base_path.exists(): + base_model_path = str(relative_base_path) + + self._load_base_hf(base_model_path, strict=strict) + return self._load_adapter(adapter_dir, strict=strict) diff --git a/xtuner/v1/module/lora_linear/__init__.py b/xtuner/v1/module/lora_linear/__init__.py index e69de29bb2..6fa7691f79 100644 --- a/xtuner/v1/module/lora_linear/__init__.py +++ b/xtuner/v1/module/lora_linear/__init__.py @@ -0,0 +1,5 @@ +from .lora_grouped_linear import LoraGroupedLinear +from .lora_linear import LoraLinear + + +__all__ = ["LoraGroupedLinear", "LoraLinear"] diff --git a/xtuner/v1/module/lora_linear/lora_grouped_linear.py b/xtuner/v1/module/lora_linear/lora_grouped_linear.py index 8827e35b33..921b44988c 100644 --- a/xtuner/v1/module/lora_linear/lora_grouped_linear.py +++ b/xtuner/v1/module/lora_linear/lora_grouped_linear.py @@ -4,13 +4,16 @@ import torch.nn as nn from torch.distributed.tensor import DTensor +from xtuner.v1.float8.float8_gmm_tile_wise import TileWiseFloat8GroupedLinear +from xtuner.v1.utils import init_params + from ..grouped_linear.moe_group_linear import GroupedLinear, build_grouped_linear class LoraGroupedLinear(nn.Module): def __init__( self, - base_layer: GroupedLinear, + base_layer: GroupedLinear | TileWiseFloat8GroupedLinear, rank: int, alpha: int, lora_dropout: float = 0.0, @@ -27,6 +30,7 @@ def __init__( self.alpha = alpha self.scale = alpha / rank self.merged = False + self.init_lora_weights = init_lora_weights self.lora_A = build_grouped_linear(self.in_features, self.rank, self.num_routed_experts, ep_mesh=self.ep_mesh) self.lora_B = build_grouped_linear(self.rank, self.out_features, self.num_routed_experts, ep_mesh=self.ep_mesh) @@ -36,31 +40,31 @@ def __init__( for p in self.base_layer.parameters(): p.requires_grad = False - if init_lora_weights: - self.reset_parameters() + self.reset_parameters() def reset_parameters(self): - if isinstance(self.lora_A.weight, DTensor): - # TODO: init DTensor - raise NotImplementedError + if self.lora_A.weight.is_meta: + return + + init_params(self.lora_A.weight, lambda weight: nn.init.kaiming_uniform_(weight, a=math.sqrt(5))) + if self.init_lora_weights: + init_params(self.lora_B.weight, nn.init.zeros_) else: - nn.init.kaiming_uniform_(self.lora_A.weight, a=math.sqrt(5)) - nn.init.zeros_(self.lora_B.weight) + init_params(self.lora_B.weight, lambda weight: nn.init.kaiming_uniform_(weight, a=math.sqrt(5))) def forward(self, x: torch.Tensor, tokens_per_expert: torch.Tensor, decoding: bool = False) -> torch.Tensor: if self.merged: return self.base_layer(x, tokens_per_expert, decoding) original_out = self.base_layer(x, tokens_per_expert, decoding) - # lora_out = self.lora_a_naive_forward(x, tokens_per_expert, decoding) - lora_out = self.lora_A(x, tokens_per_expert, decoding) - # lora_out = self.lora_b_naive_forward(lora_out, tokens_per_expert, decoding) + lora_out = self.lora_A(self.lora_dropout(x), tokens_per_expert, decoding) lora_out = self.lora_B(lora_out, tokens_per_expert, decoding) return original_out + lora_out * self.scale def lora_a_naive_forward( self, x: torch.Tensor, tokens_per_expert: torch.Tensor, decoding: bool = False ) -> torch.Tensor: - weight = self.lora_A.weight.view(-1, self.rank, self.in_features) + weight = self.lora_A.weight.to_local() if isinstance(self.lora_A.weight, DTensor) else self.lora_A.weight + weight = weight.view(-1, self.rank, self.in_features) batch_sizes = tokens_per_expert.cpu().numpy() out = [] @@ -74,7 +78,8 @@ def lora_a_naive_forward( def lora_b_naive_forward( self, x: torch.Tensor, tokens_per_expert: torch.Tensor, decoding: bool = False ) -> torch.Tensor: - weight = self.lora_B.weight.view(-1, self.out_features, self.rank) + weight = self.lora_B.weight.to_local() if isinstance(self.lora_B.weight, DTensor) else self.lora_B.weight + weight = weight.view(-1, self.out_features, self.rank) batch_sizes = tokens_per_expert.cpu().numpy() out = [] @@ -87,11 +92,87 @@ def lora_b_naive_forward( @torch.no_grad() def merge_lora(self): - raise NotImplementedError + """把 LoRA 权重合并进 base_layer.weight。 + + 对每个 expert e，计算 delta_W_e = B_e @ A_e，然后加到 base weight 上： + base_weight_e += delta_W_e * scale + 其中 base_weight_e、A_e、B_e 分别是 base_layer、lora_A、lora_B 的第 e 个 expert 的权重。 + """ + if self.merged: + return + + base_weight = self.base_layer.weight + lora_a_weight = self.lora_A.weight + lora_b_weight = self.lora_B.weight + + if isinstance(base_weight, DTensor): + # DTensor 场景：在 local tensor 上操作，每个 rank 持有部分 expert + base_local = base_weight.to_local() + a_local = lora_a_weight.to_local() + b_local = lora_b_weight.to_local() + else: + base_local = base_weight + a_local = lora_a_weight + b_local = lora_b_weight + + local_experts = a_local.shape[0] // self.rank + expected_base_rows = local_experts * self.out_features + if ( + a_local.shape != (local_experts * self.rank, self.in_features) + or b_local.shape != (local_experts * self.out_features, self.rank) + or base_local.ndim != 2 + or base_local.shape[0] < expected_base_rows + or base_local.shape[1] != self.in_features + ): + raise RuntimeError("Grouped LoRA weights must contain complete local experts before merge_lora()") + + base_view = base_local[:expected_base_rows].view(local_experts, self.out_features, self.in_features) + a_view = a_local.view(local_experts, self.rank, self.in_features) + b_view = b_local.view(local_experts, self.out_features, self.rank) + base_view.add_(torch.bmm(b_view, a_view), alpha=self.scale) + + self.merged = True @torch.no_grad() def unmerge_lora(self): - raise NotImplementedError + """从 base_layer.weight 中还原 LoRA（如果之前 merge 过）。 + + 对每个 expert e，从 base weight 中减去 delta_W_e * scale： + base_weight_e -= delta_W_e * scale + """ + if not self.merged: + return + + base_weight = self.base_layer.weight + lora_a_weight = self.lora_A.weight + lora_b_weight = self.lora_B.weight + + if isinstance(base_weight, DTensor): + base_local = base_weight.to_local() + a_local = lora_a_weight.to_local() + b_local = lora_b_weight.to_local() + else: + base_local = base_weight + a_local = lora_a_weight + b_local = lora_b_weight + + local_experts = a_local.shape[0] // self.rank + expected_base_rows = local_experts * self.out_features + if ( + a_local.shape != (local_experts * self.rank, self.in_features) + or b_local.shape != (local_experts * self.out_features, self.rank) + or base_local.ndim != 2 + or base_local.shape[0] < expected_base_rows + or base_local.shape[1] != self.in_features + ): + raise RuntimeError("Grouped LoRA weights must contain complete local experts before unmerge_lora()") + + base_view = base_local[:expected_base_rows].view(local_experts, self.out_features, self.in_features) + a_view = a_local.view(local_experts, self.rank, self.in_features) + b_view = b_local.view(local_experts, self.out_features, self.rank) + base_view.sub_(torch.bmm(b_view, a_view), alpha=self.scale) + + self.merged = False def __repr__(self) -> str: - return "lora." + super().__repr__() \ No newline at end of file + return "lora." + super().__repr__() diff --git a/xtuner/v1/module/lora_linear/lora_linear.py b/xtuner/v1/module/lora_linear/lora_linear.py index 9f2580ae37..687dae3419 100644 --- a/xtuner/v1/module/lora_linear/lora_linear.py +++ b/xtuner/v1/module/lora_linear/lora_linear.py @@ -2,6 +2,9 @@ import torch import torch.nn as nn +from torch.distributed.tensor import DTensor + +from xtuner.v1.utils import init_params from ..linear.linear import build_linear @@ -24,6 +27,7 @@ def __init__( self.alpha = alpha self.scale = alpha / rank self.merged = False + self.init_lora_weights = init_lora_weights weight = base_layer.weight dtype = weight.dtype @@ -39,12 +43,17 @@ def __init__( for p in self.base_layer.parameters(): p.requires_grad = False - if init_lora_weights: - self.reset_parameters() + self.reset_parameters() def reset_parameters(self): - nn.init.kaiming_uniform_(self.lora_A.weight, a=math.sqrt(5)) - nn.init.zeros_(self.lora_B.weight) + if self.lora_A.weight.is_meta: + return + + init_params(self.lora_A.weight, lambda weight: nn.init.kaiming_uniform_(weight, a=math.sqrt(5))) + if self.init_lora_weights: + init_params(self.lora_B.weight, nn.init.zeros_) + else: + init_params(self.lora_B.weight, lambda weight: nn.init.kaiming_uniform_(weight, a=math.sqrt(5))) def forward(self, x: torch.Tensor) -> torch.Tensor: if self.merged: @@ -60,16 +69,23 @@ def merge_lora(self): if self.merged: return - # delta_W = B @ A shape: [out, in] - delta_w = torch.matmul(self.lora_B.weight, self.lora_A.weight) - self.base_layer.weight += delta_w * self.scale + base_weight = self.base_layer.weight + lora_a_weight = self.lora_A.weight + lora_b_weight = self.lora_B.weight + + base_local = base_weight.to_local() if isinstance(base_weight, DTensor) else base_weight + a_local = lora_a_weight.to_local() if isinstance(lora_a_weight, DTensor) else lora_a_weight + b_local = lora_b_weight.to_local() if isinstance(lora_b_weight, DTensor) else lora_b_weight + + if a_local.shape != (self.rank, self.in_features) or b_local.shape != (self.out_features, self.rank): + raise RuntimeError("LoRA weights must be unsharded before merge_lora() is called") + if base_local.ndim != 2 or base_local.shape[0] < self.out_features or base_local.shape[1] != self.in_features: + raise RuntimeError("Base weight has an incompatible shape for merge_lora()") + + delta_w = torch.matmul(b_local, a_local) + base_local[: self.out_features].add_(delta_w, alpha=self.scale) self.merged = True - # 合并后 GoRA 参数可以不再训练 - for p in self.lora_A.parameters(): - p.requires_grad = False - for p in self.lora_B.parameters(): - p.requires_grad = False @torch.no_grad() def unmerge_lora(self): @@ -77,14 +93,23 @@ def unmerge_lora(self): if not self.merged: return - delta_w = torch.matmul(self.lora_B.weight, self.lora_A.weight) - self.base_layer.weight -= delta_w * self.scale + base_weight = self.base_layer.weight + lora_a_weight = self.lora_A.weight + lora_b_weight = self.lora_B.weight + + base_local = base_weight.to_local() if isinstance(base_weight, DTensor) else base_weight + a_local = lora_a_weight.to_local() if isinstance(lora_a_weight, DTensor) else lora_a_weight + b_local = lora_b_weight.to_local() if isinstance(lora_b_weight, DTensor) else lora_b_weight + + if a_local.shape != (self.rank, self.in_features) or b_local.shape != (self.out_features, self.rank): + raise RuntimeError("LoRA weights must be unsharded before unmerge_lora() is called") + if base_local.ndim != 2 or base_local.shape[0] < self.out_features or base_local.shape[1] != self.in_features: + raise RuntimeError("Base weight has an incompatible shape for unmerge_lora()") + + delta_w = torch.matmul(b_local, a_local) + base_local[: self.out_features].sub_(delta_w, alpha=self.scale) self.merged = False - for p in self.lora_A.parameters(): - p.requires_grad = True - for p in self.lora_B.parameters(): - p.requires_grad = True def extra_repr(self) -> str: return ( @@ -97,4 +122,4 @@ def extra_repr(self) -> str: ) def __repr__(self) -> str: - return "lora." + super().__repr__() \ No newline at end of file + return "lora." + super().__repr__() diff --git a/xtuner/v1/rl/trainer/worker.py b/xtuner/v1/rl/trainer/worker.py index 7890575929..332f7e97d5 100644 --- a/xtuner/v1/rl/trainer/worker.py +++ b/xtuner/v1/rl/trainer/worker.py @@ -39,7 +39,7 @@ from xtuner.v1.loss import BaseLossContext, CELossConfig, LogProbConfig from xtuner.v1.loss.ce_loss import CELossContext, LMHeadLossContext from xtuner.v1.loss.mtp_loss import MTPLossConfig, MTPLossContext -from xtuner.v1.model.adapter.lora import LoraConfig, LoraModel +from xtuner.v1.model.adapter.lora import LoraConfig from xtuner.v1.model.base import BaseModel as XtunerBaseModel from xtuner.v1.model.base import ModelItem, TransformerConfig from xtuner.v1.model.compose.base import BaseComposeConfig, BaseComposeModel @@ -136,6 +136,7 @@ class WorkerConfig(BaseModel): loss_cfg: BaseRLLossConfig lr_cfg: LRConfig fsdp_cfg: FSDPConfig + adapter_cfg: LoraConfig | None = None load_from: str | Path # TODO: 把 actor 和 ref 配置分离 optimizer_steps: int = 1 sp_size: int = 1 diff --git a/xtuner/v1/rl/weight_update/update_weighter.py b/xtuner/v1/rl/weight_update/update_weighter.py index 28134dea81..70d9f0229f 100644 --- a/xtuner/v1/rl/weight_update/update_weighter.py +++ b/xtuner/v1/rl/weight_update/update_weighter.py @@ -89,11 +89,12 @@ def update_weights(self): # 直接加载 LoRA 参数，避免 merge/unmerge 的开销和精度损失。 lora_model.merge_lora() - self._transport.update(self.weight_iterator) - - if lora_model is not None: - # 恢复 LoRA 状态，保持训练侧模型可继续训练 - lora_model.unmerge_lora() + try: + self._transport.update(self.weight_iterator) + finally: + if lora_model is not None: + # 恢复 LoRA 状态，保持训练侧模型可继续训练 + lora_model.unmerge_lora() def _set_transport(self) -> None: rollout_info = self.rollout_info diff --git a/xtuner/v1/rl/weight_update/weight_iterator.py b/xtuner/v1/rl/weight_update/weight_iterator.py index b9a2064076..f25b84ab2b 100644 --- a/xtuner/v1/rl/weight_update/weight_iterator.py +++ b/xtuner/v1/rl/weight_update/weight_iterator.py @@ -8,6 +8,7 @@ import tqdm from torch.distributed.tensor import DTensor +from xtuner.v1.model.adapter.lora import LoraModel from xtuner.v1.model.compose.base import BaseComposeConfig from xtuner.v1.model.compose.qwen3_vl import Qwen3VLForConditionalGeneration from xtuner.v1.model.moe.moe import MoE @@ -177,6 +178,8 @@ def iter_hf_batches(self, submodule=None, final_update=False): """Update the model weights.""" model = self._engine.model + if isinstance(model, LoraModel): + model = model.base_model if submodule: model = getattr(model, submodule) @@ -264,6 +267,8 @@ def iter_layer_batches(self): """Update the model weights.""" model = self._engine.model + if isinstance(model, LoraModel): + model = model.base_model DEVICE_MODULE.empty_cache() if isinstance(model.config, BaseComposeConfig): @@ -304,6 +309,8 @@ def get_params(tensor_list, name_list, save_dtype): tensor_list = [] name_list = [] for sub_name, param in layer.state_dict().items(): + if "lora_A." in sub_name or "lora_B." in sub_name: + continue if isinstance(model.config, BaseComposeConfig): saved_list.append(f"language_model.layers.{i}.{sub_name}") else: @@ -311,11 +318,12 @@ def get_params(tensor_list, name_list, save_dtype): local_tensor = param._local_tensor if isinstance(param, DTensor) else param local_tensor = local_tensor.bfloat16() load_spec = language_model.load_spec_mapping.get(f"layers.{i}.{sub_name}") + export_sub_name = sub_name.replace(".base_layer.", ".") if isinstance(model.config, BaseComposeConfig): - name = f"model.language_model.layers.{i}.{sub_name}" + name = f"model.language_model.layers.{i}.{export_sub_name}" else: - name = f"model.layers.{i}.{sub_name}" + name = f"model.layers.{i}.{export_sub_name}" if ".experts." in name and ".mlp.experts." not in name: name = name.replace(".experts.", ".mlp.experts.") @@ -330,28 +338,31 @@ def get_params(tensor_list, name_list, save_dtype): for name, param in model.state_dict().items(): if name in saved_list: continue + if "lora_A." in name or "lora_B." in name: + continue local_tensor = param._local_tensor if isinstance(param, DTensor) else param local_tensor = local_tensor.bfloat16() load_spec = model.load_spec_mapping.get(name) + export_name = name.replace(".base_layer.", ".") if isinstance(model.config, BaseComposeConfig): - if "vision_tower." in name: - name = name.replace("vision_tower.", vision_hf_prefix) - elif "multi_modal_projector." in name: - name = name.replace("multi_modal_projector.", projector_hf_prefix) - elif name == "language_model.norm.weight": - name = "model.language_model.norm.weight" - elif name == "language_model.embed_tokens.weight": - name = "model.language_model.embed_tokens.weight" - elif name == "language_model.lm_head.weight": - name = "lm_head.weight" + if "vision_tower." in export_name: + export_name = export_name.replace("vision_tower.", vision_hf_prefix) + elif "multi_modal_projector." in export_name: + export_name = export_name.replace("multi_modal_projector.", projector_hf_prefix) + elif export_name == "language_model.norm.weight": + export_name = "model.language_model.norm.weight" + elif export_name == "language_model.embed_tokens.weight": + export_name = "model.language_model.embed_tokens.weight" + elif export_name == "language_model.lm_head.weight": + export_name = "lm_head.weight" else: - if name == "norm.weight": - name = "model.norm.weight" - elif name == "embed_tokens.weight": - name = "model.embed_tokens.weight" + if export_name == "norm.weight": + export_name = "model.norm.weight" + elif export_name == "embed_tokens.weight": + export_name = "model.embed_tokens.weight" tensor_list = [(local_tensor, load_spec)] - name_list = [name] + name_list = [export_name] fsdp_unshard_tensor_list, name_list = get_params(tensor_list, name_list, dtype) state_dict = dict(zip(name_list, fsdp_unshard_tensor_list)) yield WeightUpdateBatch(state_dict) diff --git a/xtuner/v1/train/trainer.py b/xtuner/v1/train/trainer.py index f96327f6db..ea95c88fd4 100644 --- a/xtuner/v1/train/trainer.py +++ b/xtuner/v1/train/trainer.py @@ -1120,7 +1120,7 @@ def build_engine( model_config: XTunerBaseModelConfig, optim_config: OptimConfig, fsdp_config: FSDPConfig, - adapter_config: LoraConfig, + adapter_config: LoraConfig | None, load_checkpoint_path: str | Path | None, intra_layer_micro_batch: int = 1, strict: bool = True, @@ -1144,7 +1144,7 @@ def build_engine( optim_cfg=optim_config, fsdp_cfg=fsdp_config, model_cfg=model_config, - adapter_config=adapter_config, + adapter_cfg=adapter_config, intra_layer_micro_batch=intra_layer_micro_batch, ) if model_path is not None and (model_config.dcp_ignore_frozen_params or load_checkpoint_path is None):