Publish TrainerRank checkpoint slots

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]

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

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,