Record the expert LoRA layout in checkpoint metadata

Author: yueming-yuan

miles/tinker/core/input_validation.py

@@ -130,10 +130,13 @@ def validate_checkpoint_compatibility(meta: dict, record: ModelRecord, config: G
         "base_model": record.base_model,
         "lora_rank": record.lora_rank,
         "lora_alpha": record.lora_alpha,
+        "experts_shared_outer_loras": config.experts_shared_outer_loras,
         "train_attn": config.trains_attn,
         "train_mlp": config.trains_mlp,
         "train_unembed": config.trains_unembed,
     }
+    # Checkpoints written before the layout was recorded are all per-expert.
+    meta = {"experts_shared_outer_loras": False, **meta}
     for key, value in expected.items():
         if meta[key] != value:
             raise UserInputError(

miles/tinker/core/types.py

@@ -58,6 +58,7 @@ class GatewayConfig:
     max_samples_per_request: int = 64
     max_lora_rank: int = 32  # the slot capacity the server was built with (--lora-rank)
     lora_alpha: float | None = None  # None: 2 * rank
+    experts_shared_outer_loras: bool = False
     # what the server-wide adapter layout trains; create_model rejects deviations
     trains_attn: bool = True
     trains_mlp: bool = True

miles/tinker/core/utils.py

@@ -33,6 +33,7 @@ def build_checkpoint_metadata(record: ModelRecord, config: GatewayConfig) -> dic
         "base_model": record.base_model,
         "lora_rank": record.lora_rank,
         "lora_alpha": record.lora_alpha,
+        "experts_shared_outer_loras": config.experts_shared_outer_loras,
         "train_attn": config.trains_attn,
         "train_mlp": config.trains_mlp,
         "train_unembed": config.trains_unembed,

serve_tinker.py

@@ -78,6 +78,7 @@ async def serve(args, *, disposer: Disposer):
         max_tokens_per_datum=max_tokens_per_datum,
         max_tokens_per_sample=max_tokens_per_sample,
         lora_alpha=args.lora_alpha,
+        experts_shared_outer_loras=args.experts_shared_outer_loras,
         max_lora_rank=args.lora_rank,
         trains_attn="attn" in args.tinker_lora_groups,
         trains_mlp="mlp" in args.tinker_lora_groups,

tests/fast/tinker/core/test_service.py

@@ -676,7 +676,8 @@ async def test_checkpoint_meta_stores_a_digest_not_the_credential(service):
     assert (info["train_attn"], info["train_mlp"], info["train_unembed"]) == (True, True, False)
 
 
-async def test_a_checkpoint_saved_under_other_settings_does_not_load(service):
+@pytest.mark.parametrize("field,value", [("lora_alpha", 99), ("experts_shared_outer_loras", True)])
+async def test_a_checkpoint_saved_under_other_settings_does_not_load(service, field, value):
     import json
     import os
 
@@ -689,14 +690,14 @@ async def test_a_checkpoint_saved_under_other_settings_does_not_load(service):
         resolve_checkpoint_dir(service.config.checkpoint_root, model_id, "weights", "ck"), "META.json"
     )
     meta = json.loads(open(meta_path).read())
-    meta["lora_alpha"] = meta["lora_alpha"] + 1  # the same tensors would be scaled differently
+    meta[field] = value
     open(meta_path, "w").write(json.dumps(meta))
 
     loaded = service.submit(
         "tenant", "load_state", {"model_id": model_id, "seq_id": 2, "path": path, "optimizer": True}
     )
     future = await await_settled(service, "tenant", loaded)
-    assert (future.state, future.error_category) == (FAILED, "user") and "lora_alpha" in future.error
+    assert (future.state, future.error_category) == (FAILED, "user") and field in future.error
     assert not service.backend.named("load_slot")[1:], "nothing may touch the slot on a mismatch"