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..12a498ba4b --- /dev/null +++ b/tests/engine/test_dense_lora_train_engine.py @@ -0,0 +1,209 @@ +import os +import shutil +import tempfile +import time + +import parametrize +import torch +import torch.distributed as dist +from torch.optim.lr_scheduler import LambdaLR + +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 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 + + +# 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 = [] + + 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 + 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() + 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 Exception: + 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, + 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) + + torch.cuda.empty_cache() + try: + dist.destroy_process_group(pg) + except Exception: + 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 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/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..f6f460a553 --- /dev/null +++ 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 new file mode 100644 index 0000000000..d885b78554 --- /dev/null +++ b/xtuner/v1/model/adapter/lora.py @@ -0,0 +1,789 @@ +import functools +import json +import os +import re +from concurrent.futures import Future +from itertools import chain +from pathlib import Path +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 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_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() + + +class LoraConfig(PydanticBaseModel): + 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: 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) + + 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: dict[str, Any] = { + "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, 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): + 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): + _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 + + self._apply_lora() + 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() + + 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}" + ) + + 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): + 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)): + 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, + 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, 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, + 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 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)): + bias = getattr(module.base_layer, "bias", None) + if bias is not None: + bias.requires_grad = True + return + + 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 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. + + 主要用于 RL 权重更新场景：推理引擎目前只支持接收全量参数（base weight + LoRA delta）， + 不支持直接加载 LoRA 参数。因此在 weight update 前需要临时 merge，update 后再 unmerge。 + """ + 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): + """Unmerge LoRA weights from base layers. + + 与 merge_lora() 配对使用，在 RL 权重更新完成后恢复 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] + 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) + 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 + 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): + 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 a PEFT-compatible adapter checkpoint.""" + if isinstance(hf_dir, str): + hf_dir = Path(hf_dir) + hf_dir.mkdir(parents=True, exist_ok=True) + + 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() + + 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) + + if dist.is_initialized(): + torch.distributed.barrier() + + 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() + + 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}" + ) + + 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]] = [] + + 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"), + ), + ] + ) + + 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: + 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}" + ) + 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 new file mode 100644 index 0000000000..6fa7691f79 --- /dev/null +++ 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 new file mode 100644 index 0000000000..921b44988c --- /dev/null +++ b/xtuner/v1/module/lora_linear/lora_grouped_linear.py @@ -0,0 +1,178 @@ +import math + +import torch +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 | TileWiseFloat8GroupedLinear, + 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.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) + + 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 + + self.reset_parameters() + + def reset_parameters(self): + 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, 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(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.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 = [] + 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.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 = [] + 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): + """把 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): + """从 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__() 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..687dae3419 --- /dev/null +++ b/xtuner/v1/module/lora_linear/lora_linear.py @@ -0,0 +1,125 @@ +import math + +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 + + +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 + self.init_lora_weights = init_lora_weights + + 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 + + self.reset_parameters() + + def reset_parameters(self): + 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: + 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 + + 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 + + @torch.no_grad() + def unmerge_lora(self): + """从 base_layer.weight 中还原 LoRA（如果之前 merge 过）""" + 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 + + 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 + + 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__() diff --git a/xtuner/v1/rl/trainer/worker.py b/xtuner/v1/rl/trainer/worker.py index 3f4f122f8b..332f7e97d5 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 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 @@ -135,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 @@ -323,6 +325,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..70d9f0229f 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,7 +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." - self._transport.update(self.weight_iterator) + + 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() + + 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 55f236ba11..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) @@ -235,11 +238,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 @@ -255,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): @@ -295,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: @@ -302,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.") @@ -321,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 e61f46b224..ea95c88fd4 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 | None, 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_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):