[1/7][multi-lora]: utils foundation — sample/adapter types, adapter yaml config, shared helpers, CLI flags and validation

miles/rollout/generate_utils/sample_utils.py

@@ -142,6 +142,8 @@ def _merge_metadata():
             metadata=_merge_metadata(),
             generate_function_path=_merge_equal_value("generate_function_path"),
             train_metadata=_merge_equal_value("train_metadata"),
+            adapter=_merge_equal_value("adapter"),
+            reward_spec=_merge_equal_value("reward_spec"),
             routing_key=_merge_equal_value("routing_key"),
             non_generation_time=_merge_equal_value("non_generation_time"),
             spec_info=_merge_spec_info(a.spec_info, b.spec_info),

miles/utils/adapter_config.py

@@ -0,0 +1,93 @@
+"""Adapter config parsing for multi-LoRA training.
+
+``AdapterRunConfig`` carries only static, YAML-sourced configuration; the
+mutable slot is owned by the controller and exposed through ``AdapterRun``
+views.
+"""
+
+from dataclasses import dataclass, field
+from pathlib import Path
+from typing import Any
+
+import yaml
+
+
+@dataclass(frozen=True)
+class AdapterRunConfig:
+
+    data: str
+
+    # resolves them to CLI defaults if None (--lora-rank / --lora-alpha) on register.
+    rank: int | None = None
+    alpha: int | None = None
+
+    # Prompt groups consumed per optimizer step for this adapter (group units,
+    # like --rollout-batch-size, which it defaults to). The samples-per-step
+    # analog of --global-batch-size is derived: adapter_global_batch_size =
+    # rollout_batch_size * n_samples_per_prompt.
+    rollout_batch_size: int | None = None
+    n_samples_per_prompt: int | None = None
+
+    save: str | Path | None = None
+
+    input_key: str = "text"
+    label_key: str | None = None
+    metadata_key: str | None = None
+
+    rm_type: str | None = None
+    custom_rm_path: str | None = None
+
+    # Stop after N optimizer steps; derived from num_epoch (default 1) when absent.
+    num_step: int | None = None
+    num_epoch: int | None = None
+
+    metadata: dict[str, Any] = field(default_factory=dict)
+
+    @property
+    def adapter_global_batch_size(self) -> int:
+        """Samples per optimizer step (per-adapter analog of --global-batch-size)."""
+        assert self.rollout_batch_size is not None and self.n_samples_per_prompt is not None
+        return self.rollout_batch_size * self.n_samples_per_prompt
+
+
+@dataclass(frozen=True)
+class AdapterRun:
+    """Read-only join view of a run's static config and current slot."""
+
+    name: str
+    config: AdapterRunConfig
+    slot: int
+    version: int = 0
+    step: int = 0
+    # Committed prompt groups accumulated toward the current optimizer step.
+    accumulated_groups: int = 0
+    # Unique per registration (see AdapterRecord.registration_id): lets the
+    # rollout worker tell a re-registered name apart from the previous tenant.
+    registration_id: str = ""
+
+
+def parse_adapter_run_yaml(path: Path) -> AdapterRunConfig:
+    """Parse a single adapter.yaml file.
+
+    ``rank``, ``alpha`` and ``save`` are optional in the YAML; when absent the
+    caller (e.g. the multi-LoRA controller) is responsible for resolving them.
+    """
+    with open(path) as f:
+        raw = yaml.safe_load(f)
+
+    return AdapterRunConfig(
+        rank=raw.get("rank"),
+        alpha=raw.get("alpha"),
+        data=raw["data"],
+        rollout_batch_size=raw.get("rollout_batch_size"),
+        n_samples_per_prompt=raw.get("n_samples_per_prompt"),
+        save=Path(raw["save"]) if raw.get("save", None) else None,
+        input_key=raw.get("input_key", "text"),
+        label_key=raw.get("label_key"),
+        metadata_key=raw.get("metadata_key"),
+        rm_type=raw.get("rm_type"),
+        custom_rm_path=raw.get("custom_rm_path"),
+        num_step=raw.get("num_step"),
+        num_epoch=raw.get("num_epoch"),
+        metadata=raw.get("metadata") or {},
+    )

miles/utils/arguments.py

@@ -538,7 +538,7 @@ def add_rollout_arguments(parser):
                 default=512 * 1024**2,
                 help=(
                     "buffer size for update weight, in bytes. "
-                    "This is used for updating weights by chunk and should be useful for MoE models."
+                    "This is used for updating weights by batch and should be useful for MoE models."
                 ),
             )
             parser.add_argument(
@@ -1400,6 +1400,84 @@ def add_lora_arguments(parser):
                 "down lora_B shared across experts, expert_dim=1). Matches SGLang "
                 "PR #21466's experts_shared_outer_loras=True serving contract.",
             )
+            parser.add_argument(
+                "--multi-lora-n-adapters",
+                type=int,
+                default=0,
+                help="Maximum number of concurrent adapter slots for multi-LoRA. Set to 0 to disable multi-LoRA (default: 0)",
+            )
+            parser.add_argument(
+                "--multi-lora-adapter",
+                nargs=2,
+                action="append",
+                type=str,
+                dest="multi_lora_adapters",
+                default=[],
+            )
+            parser.add_argument(
+                "--multi-lora-idle-poll-s",
+                type=float,
+                default=5.0,
+                help="When no adapter is RUNNING, the trainer polls for new registrations every this many seconds (default: 5.0)",
+            )
+            parser.add_argument(
+                "--multi-lora-http-server-path",
+                type=str,
+                default=None,
+                help=(
+                    "Dotted path to a MultiLoRAHTTPServer subclass to use for the multi-LoRA "
+                    "controller's HTTP server (default: MultiLoRAHTTPServer)"
+                ),
+            )
+            parser.add_argument(
+                "--multi-lora-backend-path",
+                type=str,
+                default=None,
+                help=(
+                    "Dotted path to a MultiLoRABackend subclass for the multi-LoRA controller, "
+                    "e.g. to add custom adapter validation via validate_adapter (default: MultiLoRABackend)"
+                ),
+            )
+            parser.add_argument(
+                "--multi-lora-api-port",
+                type=int,
+                default=8068,
+                help="Port for the multi-LoRA controller's control-plane API, served from the head node (default: 8068)",
+            )
+            parser.add_argument(
+                "--multi-lora-disable-service-mode",
+                action="store_false",
+                dest="multi_lora_service_mode",
+                help="Disable service mode. By default, the trainer waits indefinitely for new adapters. With this flag, it exits after all adapters have been processed.",
+            )
+            parser.add_argument(
+                "--multi-lora-max-adapter-global-batch-size",
+                type=int,
+                default=None,
+                help=(
+                    "Registration-time upper bound on an adapter's samples per optimizer "
+                    "step (rollout_batch_size x n_samples_per_prompt). Defaults to 4x "
+                    "--global-batch-size."
+                ),
+            )
+            parser.add_argument(
+                "--multi-lora-max-coalesce-wait-s",
+                type=float,
+                default=0.5,
+                help=(
+                    "Maximum time ready groups wait for the batch to fill toward "
+                    "--global-batch-size before training starts on what is ready (default: 0.5)."
+                ),
+            )
+            parser.add_argument(
+                "--multi-lora-max-empty-wait-s",
+                type=float,
+                default=30.0,
+                help=(
+                    "How long a generate call waits for the first poppable group before "
+                    "failing with an empty-batch timeout (default: 30)."
+                ),
+            )
             return parser
 
         def add_router_arguments(parser):
@@ -2532,6 +2610,12 @@ def miles_validate_args(args):
                 "shared-outer" if args.experts_shared_outer_loras else "per-expert",
             )
 
+    # Sets args.multi_lora, then validates/defaults the multi-LoRA arg surface
+    # (adapter configs themselves are loaded later by the controller).
+    from miles.utils.multi_lora import validate_multi_lora_args
+
+    validate_multi_lora_args(args)
+
     assert not (args.kl_coef != 0 and args.kl_loss_coef != 0), "Only one of kl_coef and kl_loss_coef can be set"
 
     if args.advantage_estimator in ["reinforce_plus_plus", "reinforce_plus_plus_baseline"]:
@@ -2700,7 +2784,9 @@ def miles_validate_args(args):
             )
         args.global_batch_size = global_batch_size
 
-    if args.n_samples_per_prompt == 1:
+    # Multi-LoRA adapters carry their own n_samples_per_prompt; the per-group
+    # normalization path already skips std for singleton groups.
+    if args.n_samples_per_prompt == 1 and not args.multi_lora:
         args.grpo_std_normalization = False
         logger.info("n_samples_per_prompt is set to 1, grpo_std_normalization will be set to False.")