From da98a1bf49f767e6ffa1c6ad0893ed080baac492 Mon Sep 17 00:00:00 2001 From: Brad Hilton Date: Sat, 11 Jul 2026 01:30:31 +0000 Subject: [PATCH] feat: publish TrainerRank checkpoint slots --- src/art/megatron/lora.py | 39 ++++-- src/art/megatron/weights/lora_publish.py | 23 ++-- src/art/trainer_rank/__init__.py | 114 +++++++++++++++++- .../megatron/lora/test_lora_disk_codecs.py | 48 +++++++- tests/unit/test_trainer_rank_validation.py | 82 +++++++++++++ 5 files changed, 285 insertions(+), 21 deletions(-) diff --git a/src/art/megatron/lora.py b/src/art/megatron/lora.py index d0951fda9..e6b29cb97 100644 --- a/src/art/megatron/lora.py +++ b/src/art/megatron/lora.py @@ -738,17 +738,27 @@ def _manifest_for_param(self, param: torch.nn.Parameter) -> dict[str, Any]: manifest["component_sizes"] = component_sizes return manifest - def _lora_params(self) -> list[tuple[str, torch.nn.Parameter]]: + def _lora_params( + self, ref: LoRASlotRef | None = None + ) -> list[tuple[str, torch.nn.Parameter]]: + if ref is not None: + slot = self._slot(ref) + if slot is None: + return [] + return [ + ("lora_A.weight", slot.A_T), + ("lora_B.weight", slot.B_T), + ] return [ ("lora_A.weight", self.A_T), ("lora_B.weight", self.B_T), ] def _export_items( - self, + self, ref: LoRASlotRef | None = None ) -> list[tuple[str, torch.nn.Parameter, int | None]]: export_items: list[tuple[str, torch.nn.Parameter, int | None]] = [] - for key, param in self._lora_params(): + for key, param in self._lora_params(ref): if not self._should_export_parameter(param): continue if self.num_local_experts > 1: @@ -759,15 +769,19 @@ def _export_items( export_items.append((f"{self.adapter_model_prefix}.{key}", param, None)) return export_items - def sharded_lora_manifest(self) -> dict[str, dict[str, Any]]: + def sharded_lora_manifest( + self, ref: LoRASlotRef | None = None + ) -> dict[str, dict[str, Any]]: return { key: self._manifest_for_param(param) - for key, param, _expert in self._export_items() + for key, param, _expert in self._export_items(ref) } - def sharded_lora_state_dict(self) -> dict[str, torch.Tensor]: + def sharded_lora_state_dict( + self, ref: LoRASlotRef | None = None + ) -> dict[str, torch.Tensor]: state: dict[str, torch.Tensor] = {} - for key, param, expert in self._export_items(): + for key, param, expert in self._export_items(ref): state[key] = param.data[expert].T if expert is not None else param.data.T return state @@ -822,8 +836,12 @@ def forward( class LoRAPublishPlanner: - def __init__(self, model_chunks: Sequence[torch.nn.Module]) -> None: - self.templates = tuple(self._collect_templates(model_chunks)) + def __init__( + self, + model_chunks: Sequence[torch.nn.Module], + slot_ref: LoRASlotRef | None = None, + ) -> None: + self.templates = tuple(self._collect_templates(model_chunks, slot_ref)) def global_metadata( self, @@ -846,13 +864,14 @@ def global_metadata( @staticmethod def _collect_templates( model_chunks: Sequence[torch.nn.Module], + slot_ref: LoRASlotRef | None = None, ) -> list[_LoraPublishTemplate]: templates: list[_LoraPublishTemplate] = [] for chunk in model_chunks: for module in chunk.modules(): if not isinstance(module, LoRA): continue - for suffix, param in module._lora_params(): + for suffix, param in module._lora_params(slot_ref): if not module._should_export_parameter(param): continue sharded = bool(param.lora_tp_sharded) # type: ignore[attr-defined] diff --git a/src/art/megatron/weights/lora_publish.py b/src/art/megatron/weights/lora_publish.py index f7b9c5427..27db9a9b3 100644 --- a/src/art/megatron/weights/lora_publish.py +++ b/src/art/megatron/weights/lora_publish.py @@ -7,6 +7,7 @@ LoRA, LoRAPublishPlanner, LoraShardMeta, + LoRASlotRef, _block_for_key, _dtype_name, ) @@ -98,10 +99,11 @@ def _packed_expert_slot( def _uses_packed_expert_publish( module: LoRA, groups: Sequence[ExpertPackedLoraGroup], + slot_ref: LoRASlotRef | None = None, ) -> bool: if module.num_local_experts <= 1: return False - params = tuple(module._lora_params()) + params = tuple(module._lora_params(slot_ref)) return bool(params) and all( _packed_expert_slot(module.adapter_model_prefix, suffix, groups) is not None for suffix, _param in params @@ -114,18 +116,19 @@ def collect_local_lora_entries( *, owner_rank: int, packed_expert_groups: Sequence[ExpertPackedLoraGroup] = (), + slot_ref: LoRASlotRef | None = None, ) -> tuple[dict[str, torch.Tensor], list[LoraShardMeta]]: local_tensors: dict[str, torch.Tensor] = {} local_manifest: dict[str, dict[str, Any]] = {} for module in iter_lora_modules(model_chunks): - if _uses_packed_expert_publish(module, packed_expert_groups): + if _uses_packed_expert_publish(module, packed_expert_groups, slot_ref): continue - for key, value in module.sharded_lora_state_dict().items(): + for key, value in module.sharded_lora_state_dict(slot_ref).items(): target_dtype = ( adapter_model[key].dtype if key in adapter_model else value.dtype ) local_tensors[key] = value.to(target_dtype).contiguous() - local_manifest.update(module.sharded_lora_manifest()) + local_manifest.update(module.sharded_lora_manifest(slot_ref)) if set(local_tensors) != set(local_manifest): raise RuntimeError( @@ -153,15 +156,16 @@ def collect_local_packed_expert_entries( *, owner_rank: int, packed_expert_groups: Sequence[ExpertPackedLoraGroup], + slot_ref: LoRASlotRef | None = None, ) -> tuple[dict[str, torch.Tensor], list[PackedExpertShardMeta]]: local_tensors: dict[str, torch.Tensor] = {} metadata: list[PackedExpertShardMeta] = [] for module in iter_lora_modules(model_chunks): - if not _uses_packed_expert_publish(module, packed_expert_groups): + if not _uses_packed_expert_publish(module, packed_expert_groups, slot_ref): continue expert_start = int(module._expert_offset) expert_count = int(module.num_local_experts) - for suffix, param in module._lora_params(): + for suffix, param in module._lora_params(slot_ref): slot_match = _packed_expert_slot( module.adapter_model_prefix, suffix, @@ -673,6 +677,7 @@ def build_vllm_lora_tensors_from_model( adapter_config: dict[str, Any], rank: int, world_size: int, + slot_ref: LoRASlotRef | None = None, ) -> tuple[dict[str, torch.Tensor], dict[str, Any]] | None: actual_rank, device = _rank_and_device() if _distributed_ready(): @@ -690,18 +695,20 @@ def build_vllm_lora_tensors_from_model( ) rank = 0 packed_expert_groups = tuple(handler.expert_packed_lora_groups()) - planner = LoRAPublishPlanner(model) + planner = LoRAPublishPlanner(model, slot_ref) local_tensors, local_metadata = collect_local_lora_entries( model, adapter_model, owner_rank=rank, packed_expert_groups=packed_expert_groups, + slot_ref=slot_ref, ) local_packed_tensors, local_packed_metadata = collect_local_packed_expert_entries( model, adapter_model, owner_rank=rank, packed_expert_groups=packed_expert_groups, + slot_ref=slot_ref, ) all_packed_metadata = ( _global_packed_expert_metadata(planner, adapter_model, packed_expert_groups) @@ -751,6 +758,7 @@ def save_vllm_lora_from_model( output_dir: str, rank: int, world_size: int, + slot_ref: LoRASlotRef | None = None, ) -> None: result = build_vllm_lora_tensors_from_model( model=model, @@ -759,6 +767,7 @@ def save_vllm_lora_from_model( adapter_config=adapter_config, rank=rank, world_size=world_size, + slot_ref=slot_ref, ) if result is None: return diff --git a/src/art/trainer_rank/__init__.py b/src/art/trainer_rank/__init__.py index b22483d43..f7b51bd44 100644 --- a/src/art/trainer_rank/__init__.py +++ b/src/art/trainer_rank/__init__.py @@ -6,6 +6,7 @@ Mapping, Sequence, ) +from copy import deepcopy from dataclasses import dataclass import os from typing import ( @@ -544,6 +545,7 @@ def __init__( self._checkpoint_slot_params_by_name: dict[ str, tuple[torch.nn.Parameter, ...] ] = {} + self._checkpoint_slot_adapter_configs: dict[str, dict[str, Any]] = {} self._pending_slot_graphs: dict[ LoRASlotRef, list[weakref.ReferenceType[torch.Tensor]] ] = {} @@ -592,19 +594,37 @@ def load_checkpoint_slot( *, optimizer_state: Mapping[str, object] | None = None, alpha: float | None = None, + adapter_config: Mapping[str, object] | None = None, ) -> int: + config = self._validate_checkpoint_slot_adapter_config( + name, adapter_config, alpha=alpha + ) loaded = self._load_slot( - "checkpoint", name, adapter_model, trainable=True, alpha=alpha + "checkpoint", + name, + adapter_model, + trainable=True, + alpha=alpha if config is None else float(config["lora_alpha"]), ) - self._checkpoint_slot_params_by_name[name] = ( - self._validate_dynamic_slot_consistency("checkpoint", name, loaded) + slot_params = self._validate_dynamic_slot_consistency( + "checkpoint", name, loaded ) + if config is not None: + self._validate_loaded_checkpoint_slot_config(name, config) + self._checkpoint_slot_params_by_name[name] = slot_params if optimizer_state is None: self._dynamic_optimizers.pop(name, None) else: self._dynamic_optimizers[name] = self._restore_dynamic_optimizer( name, optimizer_state ) + configs = getattr(self, "_checkpoint_slot_adapter_configs", None) + if configs is None: + configs = self._checkpoint_slot_adapter_configs = {} + if config is None: + configs.pop(name, None) + else: + configs[name] = config return loaded def checkpoint_slot_optimizer_state(self, name: str) -> dict[str, object] | None: @@ -622,6 +642,94 @@ def checkpoint_slot_optimizer_state(self, name: str) -> dict[str, object] | None "optimizer": _state_to_cpu(dynamic.optimizer.state_dict()), } + def save_checkpoint_slot_lora(self, name: str, output_dir: str) -> None: + """Collectively publish a trained checkpoint slot as a vLLM LoRA.""" + known = name in self._checkpoint_slot_params_by_name + if dist.is_available() and dist.is_initialized(): + gathered: list[tuple[str, bool] | None] = [None] * dist.get_world_size() + dist.all_gather_object(gathered, (name, known)) + if any(state != (name, True) for state in gathered): + raise ValueError( + "Checkpoint slot publish requires the same loaded name on all " + f"ranks; got {gathered}" + ) + if not known: + raise ValueError(f"Unknown checkpoint slot: {name!r}") + config = getattr(self, "_checkpoint_slot_adapter_configs", {}).get(name) + if config is None: + raise TrainerRankSlotStateError( + f"Checkpoint slot {name!r} was loaded without adapter_config; " + "reload it with adapter_config=... before publishing." + ) + from art.megatron.weights.lora_publish import save_vllm_lora_from_model + + save_vllm_lora_from_model( + model=self.runtime.model, + adapter_model={}, + handler=self.runtime.model_support_handler, + adapter_config=config, + output_dir=output_dir, + rank=self.runtime.rank, + world_size=self.runtime.world_size, + slot_ref=self._slot_ref("checkpoint", name), + ) + + def _validate_checkpoint_slot_adapter_config( + self, + name: str, + adapter_config: Mapping[str, object] | None, + *, + alpha: float | None, + ) -> dict[str, Any] | None: + config = ( + None + if adapter_config is None + else cast(dict[str, Any], deepcopy(dict(adapter_config))) + ) + if dist.is_available() and dist.is_initialized(): + gathered: list[dict[str, Any] | None] = [None] * dist.get_world_size() + dist.all_gather_object(gathered, config) + if any(value != config for value in gathered): + raise ValueError( + f"Adapter config for checkpoint slot {name!r} differs across ranks" + ) + if config is None: + return None + required = {"base_model_name_or_path", "r", "lora_alpha", "target_modules"} + if missing := sorted(required - config.keys()): + raise ValueError( + f"Adapter config for checkpoint slot {name!r} is missing {missing}" + ) + if int(config["r"]) < 1: + raise ValueError("adapter_config['r'] must be >= 1") + config_alpha = float(config["lora_alpha"]) + if alpha is not None and float(alpha) != config_alpha: + raise ValueError( + f"alpha={alpha} conflicts with adapter_config lora_alpha={config_alpha}" + ) + return config + + def _validate_loaded_checkpoint_slot_config( + self, name: str, config: Mapping[str, Any] + ) -> None: + from art.megatron.lora import LoRA + + ref = self._slot_ref("checkpoint", name) + slots = [ + slot + for chunk in self.runtime.model + for module in chunk.modules() + if isinstance(module, LoRA) + if (slot := module._slot(ref)) is not None + ] + expected = (int(config["r"]), float(config["lora_alpha"])) + actual = {(slot.rank, slot.alpha) for slot in slots} + if actual != {expected}: + raise ValueError( + f"Adapter config for checkpoint slot {name!r} declares " + f"rank/alpha={expected}, loaded weights use {sorted(actual)}" + ) + def load_lora_slot( self, name: str, diff --git a/tests/integration/megatron/lora/test_lora_disk_codecs.py b/tests/integration/megatron/lora/test_lora_disk_codecs.py index c317c52b0..7c271d1b6 100644 --- a/tests/integration/megatron/lora/test_lora_disk_codecs.py +++ b/tests/integration/megatron/lora/test_lora_disk_codecs.py @@ -5,6 +5,7 @@ import shutil import subprocess import sys +from types import SimpleNamespace from typing import Any, cast import pytest @@ -14,7 +15,7 @@ pytest.importorskip("megatron.bridge.models.gpt_provider") from art.megatron import lora as lora_module -from art.megatron.lora import LoRA, LoRAParallelSpec, LoRAPublishPlanner +from art.megatron.lora import LoRA, LoRAParallelSpec, LoRAPublishPlanner, LoRASlotRef from art.megatron.model_support.handlers import ( DEFAULT_DENSE_HANDLER, GPT_OSS_MOE_HANDLER, @@ -33,6 +34,7 @@ merge_sharded_adapter_entries, save_vllm_lora_from_model, ) +from art.trainer_rank import TrainerRank from art.utils.convert_moe_lora import convert_checkpoint_if_needed REPO_ROOT = Path(__file__).parents[4] @@ -1486,9 +1488,45 @@ def test_save_vllm_lora_from_model_writes_single_vllm_checkpoint(tmp_path: Path) _assert_tensors_equal(roundtrip, full) +def test_trainer_rank_publishes_named_checkpoint_slot_without_mutating_base( + tmp_path: Path, +): + prefix = "base_model.model.model.layers.0.self_attn.q_proj" + lora = LoRA(prefix, 3, 4, 2, 2, torch.float32, torch.device("cpu")) + baseline = (lora.A_T.detach().clone(), lora.B_T.detach().clone()) + adapter = { + f"{prefix}.lora_A.weight": torch.arange(6, dtype=torch.float32).reshape(2, 3), + f"{prefix}.lora_B.weight": torch.arange(8, dtype=torch.float32).reshape(4, 2), + } + trainer = TrainerRank.__new__(TrainerRank) + trainer.runtime = SimpleNamespace( + model=[lora], + model_support_handler=DEFAULT_DENSE_HANDLER, + rank=0, + world_size=1, + ) + trainer._slot_stack = [] + trainer._pending_slot_graphs = {} + trainer._dynamic_optimizers = {} + trainer._checkpoint_slot_params_by_name = {} + trainer._checkpoint_slot_adapter_configs = {} + config = _config("Qwen/Qwen3-8B", rank=2, alpha=2) + assert trainer.load_checkpoint_slot("student", adapter, adapter_config=config) == 1 + output_dir = tmp_path / "checkpoint" + + trainer.save_checkpoint_slot_lora("student", str(output_dir)) + + _assert_tensors_equal(load_file(output_dir / "adapter_model.safetensors"), adapter) + assert json.loads((output_dir / "adapter_config.json").read_text()) == config + assert torch.equal(lora.A_T, baseline[0]) + assert torch.equal(lora.B_T, baseline[1]) + + +@pytest.mark.parametrize("dynamic_slot", [False, True]) def test_direct_qwen35_packed_expert_publish_matches_old_vllm_exactly( tmp_path: Path, monkeypatch, + dynamic_slot: bool, ): monkeypatch.setattr(lora_module.ps, "get_expert_model_parallel_rank", lambda: 0) monkeypatch.setattr(lora_module.ps, "get_expert_data_parallel_rank", lambda: 0) @@ -1554,6 +1592,13 @@ def test_direct_qwen35_packed_expert_publish_matches_old_vllm_exactly( down_lora.B_T.data[expert].copy_(tensors["down_proj.lora_B.weight"].T) offset += 1000 + slot_ref = LoRASlotRef("checkpoint", "student") if dynamic_slot else None + if slot_ref is not None: + assert gate_up_lora.load_lora_slot( + slot_ref, full, alpha=rank, requires_grad=True + ) + assert down_lora.load_lora_slot(slot_ref, full, alpha=rank, requires_grad=True) + adapter_config = _config("Qwen/Qwen3.5-35B-A3B", rank=rank, alpha=rank) old_dir = tmp_path / "old" current_dir = tmp_path / "current" @@ -1570,6 +1615,7 @@ def test_direct_qwen35_packed_expert_publish_matches_old_vllm_exactly( output_dir=str(current_dir), rank=0, world_size=1, + slot_ref=slot_ref, ) _assert_tensors_equal( diff --git a/tests/unit/test_trainer_rank_validation.py b/tests/unit/test_trainer_rank_validation.py index b0cff972f..5bf518f11 100644 --- a/tests/unit/test_trainer_rank_validation.py +++ b/tests/unit/test_trainer_rank_validation.py @@ -222,6 +222,88 @@ def test_trainer_rank_normalizes_adapter_tensors_to_installed_site() -> None: assert all(tensor.dtype == torch.bfloat16 for tensor in normalized.values()) +def test_checkpoint_slot_adapter_config_is_validated_and_copied() -> None: + trainer = TrainerRank(_runtime()) # type: ignore[arg-type] + config = { + "base_model_name_or_path": "Qwen/Qwen3-8B", + "r": 8, + "lora_alpha": 16, + "target_modules": ["q_proj"], + } + + retained = trainer._validate_checkpoint_slot_adapter_config( + "student", config, alpha=16 + ) + + assert retained == config + config["target_modules"].append("v_proj") # type: ignore[union-attr] + assert retained is not None + assert retained["target_modules"] == ["q_proj"] + with pytest.raises(ValueError, match="conflicts"): + trainer._validate_checkpoint_slot_adapter_config("student", config, alpha=32) + with pytest.raises(ValueError, match="missing"): + trainer._validate_checkpoint_slot_adapter_config( + "student", {"r": 8}, alpha=None + ) + + +def test_checkpoint_slot_adapter_config_rejects_cross_rank_mismatch( + monkeypatch: pytest.MonkeyPatch, +) -> None: + trainer = TrainerRank(_runtime()) # type: ignore[arg-type] + monkeypatch.setattr("art.trainer_rank.dist.is_initialized", lambda: True) + monkeypatch.setattr("art.trainer_rank.dist.get_world_size", lambda: 2) + + def gather(output: list[object], value: object) -> None: + output[:] = [value, {"different": True}] + + monkeypatch.setattr("art.trainer_rank.dist.all_gather_object", gather) + + with pytest.raises(ValueError, match="differs across ranks"): + trainer._validate_checkpoint_slot_adapter_config("student", None, alpha=None) + + +def test_load_checkpoint_slot_retains_config_and_uses_its_alpha( + monkeypatch: pytest.MonkeyPatch, +) -> None: + trainer = TrainerRank(_runtime()) # type: ignore[arg-type] + seen: dict[str, object] = {} + monkeypatch.setattr( + trainer, + "_load_slot", + lambda *_args, **kwargs: seen.update(kwargs) or 1, + ) + monkeypatch.setattr(trainer, "_validate_dynamic_slot_consistency", lambda *_: ()) + monkeypatch.setattr( + trainer, "_validate_loaded_checkpoint_slot_config", lambda *_: None + ) + config = { + "base_model_name_or_path": "Qwen/Qwen3-8B", + "r": 8, + "lora_alpha": 16, + "target_modules": ["q_proj"], + } + + trainer.load_checkpoint_slot("student", {}, adapter_config=config) + + assert seen["alpha"] == 16 + assert trainer._checkpoint_slot_adapter_configs["student"] == config + trainer.load_checkpoint_slot("student", {}, alpha=7) + assert seen["alpha"] == 7 + assert "student" not in trainer._checkpoint_slot_adapter_configs + + +def test_checkpoint_slot_publish_requires_retained_adapter_config() -> None: + trainer = TrainerRank(_runtime()) # type: ignore[arg-type] + with pytest.raises(ValueError, match="Unknown checkpoint slot"): + trainer.save_checkpoint_slot_lora("missing", "/unused") + + trainer._checkpoint_slot_params_by_name["student"] = () + + with pytest.raises(TrainerRankSlotStateError, match="adapter_config"): + trainer.save_checkpoint_slot_lora("student", "/unused") + + def test_trainer_rank_default_forward_uses_explicit_base_slot() -> None: trainer = TrainerRank(_runtime()) # type: ignore[arg-type]