[multi-lora v2 1/7] LoRA model extension and slot actor

miles/backends/megatron_utils/actor.py

@@ -13,7 +13,10 @@
 from torch_memory_saver import torch_memory_saver
 
 from miles.backends.megatron_utils.ft.types import TrainStepOutput
+from miles.backends.megatron_utils.lora.utils import build_lora_sync_config, is_lora_enabled, lora_rollout_enabled
 from miles.backends.megatron_utils.rematerialize_utils import build_main_cast_context
+from miles.backends.megatron_utils.update_weight.hf_weight_iterator import get_hf_weight_iterator
+from miles.backends.training_utils.weight_update.updater import WeightUpdater
 from miles.dashboard import hooks as dashboard_hooks
 from miles.ray.specs.train import compute_trainer_pool_id
 from miles.ray.train_actor import TrainRayActor
@@ -55,7 +58,6 @@
 from .ft.in_memory_checkpoint import InMemoryCheckpointManager
 from .ft.indep_dp import reconfigure_indep_dp_group
 from .initialize import init, is_first_replica_megatron_main_rank
-from .lora.utils import is_lora_enabled, lora_rollout_enabled
 from .model import TrainStepOutcome, forward_only, initialize_model_and_optimizer, save, train
 from .named_weights import named_params_and_buffers
 from .parallel import verify_megatron_parallel_state
@@ -246,29 +248,7 @@ def init(
         if self.args.vocab_size is None:
             self.args.vocab_size = self.tokenizer.vocab_size
 
-        from miles.backends.training_utils.weight_update.updater import WeightUpdater
-
-        from .lora.utils import build_lora_sync_config
-        from .update_weight.hf_weight_iterator import get_hf_weight_iterator
-
-        is_lora = lora_rollout_enabled(args)
-        uses_colocate_protocol = self.args.colocate
-        if is_lora and not uses_colocate_protocol:
-            assert args.megatron_to_hf_mode == "bridge", (
-                "LoRA weight sync over distributed engines requires "
-                f"--megatron-to-hf-mode bridge (got {args.megatron_to_hf_mode!r})."
-            )
-        self.weight_updater = WeightUpdater(
-            self.args,
-            self.model,
-            weights_getter=self._get_actor_weights,
-            model_name=type(self.hf_config).__name__.lower() if self.args.model_name is None else self.args.model_name,
-            quantization_config=getattr(self.hf_config, "quantization_config", None),
-            iterator_factory=get_hf_weight_iterator,
-            parallel_state=get_parallel_state(),
-            is_lora=is_lora,
-            lora_sync_config=build_lora_sync_config(self.args) if is_lora else None,
-        )
+        self._init_training_state()
 
         # empty cache after initialization
         clear_memory()
@@ -287,6 +267,29 @@ def init(
 
         return start_rollout_id
 
+    def _init_training_state(self) -> None:
+        args = self.args
+        is_lora = lora_rollout_enabled(args)
+        uses_colocate_protocol = self.args.colocate
+        if is_lora and not uses_colocate_protocol:
+            assert args.megatron_to_hf_mode == "bridge", (
+                "LoRA weight sync over distributed engines requires "
+                f"--megatron-to-hf-mode bridge (got {args.megatron_to_hf_mode!r})."
+            )
+        model_name = type(self.hf_config).__name__.lower() if args.model_name is None else args.model_name
+        quantization_config = getattr(self.hf_config, "quantization_config", None)
+        self.weight_updater = WeightUpdater(
+            args,
+            self.model,
+            weights_getter=self._get_actor_weights,
+            model_name=model_name,
+            quantization_config=quantization_config,
+            iterator_factory=get_hf_weight_iterator,
+            parallel_state=get_parallel_state(),
+            is_lora=is_lora,
+            lora_sync_config=build_lora_sync_config(args) if is_lora else None,
+        )
+
     def _clear_quantized_weight_workspaces(self) -> None:
         if not (
             self.args.clear_quantized_weight_workspaces_on_offload

miles/backends/megatron_utils/lora/actor.py

@@ -0,0 +1,49 @@
+from contextlib import ExitStack
+
+import ray
+
+from miles.backends.megatron_utils.actor import MegatronTrainRayActor
+from miles.backends.megatron_utils.lora import model as lora_model
+from miles.backends.megatron_utils.lora.optimizer import SlotOptimizer
+from miles.backends.training_utils.data import get_data_iterator, get_rollout_data
+from miles.backends.training_utils.parallel import get_parallel_state
+from miles.utils.ray_utils import Box
+from miles.utils.tracking_utils.structured_log import with_logs
+
+
+class MultiLoRATrainRayActor(MegatronTrainRayActor):
+    def _init_training_state(self) -> None:
+        self.slot_optimizers: dict[int, SlotOptimizer] = {}
+
+    @with_logs
+    def forward_backward(self, batch_id: int, rollout_data_ref: Box) -> dict:
+        self._heartbeat.bump()
+        with ExitStack() as stack:
+            rollout_data, store_get_result = get_rollout_data(self.args, rollout_data_ref)
+            stack.enter_context(store_get_result)
+            return lora_model.run_forward_backward(self.args, batch_id, self.model, rollout_data)
+
+    @with_logs
+    def optim_step(self, adam_params_by_slot: dict[int, dict]) -> dict[int, float]:
+        self._heartbeat.bump()
+        return lora_model.optim_step(self.args, self.slot_optimizers, adam_params_by_slot)
+
+    @with_logs
+    def forward_only_logprobs(self, batch_id: int, rollout_data_ref: Box) -> Box | None:
+        self._heartbeat.bump()
+        with ExitStack() as stack:
+            rollout_data, store_get_result = get_rollout_data(self.args, rollout_data_ref)
+            stack.enter_context(store_get_result)
+            data_iterator, num_microbatches = get_data_iterator(self.args, self.model, rollout_data)
+            outputs = self.compute_log_prob(data_iterator, num_microbatches, rollout_id=batch_id)
+        if not get_parallel_state().is_pp_last_stage:
+            return None
+        return Box(ray.put({key: [t.cpu() for t in tensors] for key, tensors in outputs.items()}))
+
+    @with_logs
+    def load_slot(self, slot: int, rank: int, alpha: float) -> None:
+        self.slot_optimizers[slot] = lora_model.load_slot(self.args, self.model, slot, rank, alpha)
+
+    @with_logs
+    def unload_slot(self, slot: int) -> None:
+        lora_model.unload_slot(self.model, self.slot_optimizers.pop(slot))

miles/backends/megatron_utils/lora/model.py

@@ -0,0 +1,73 @@
+"""Multi-LoRA slot operations extending the shared Megatron model execution."""
+
+from argparse import Namespace
+from collections.abc import Sequence
+
+from megatron.core.distributed import DistributedDataParallel as DDP
+
+from miles.backends.megatron_utils.lora.optimizer import (
+    SlotOptimizer,
+    reset_grad_metadata_keep_grads,
+    step_slot_optimizers,
+)
+from miles.backends.megatron_utils.model import run_forward_backward_pass, setup_train_iteration_config
+from miles.backends.training_utils.data import get_data_iterator
+from miles.backends.training_utils.log_utils import aggregate_train_losses
+from miles.backends.training_utils.parallel import get_parallel_state
+from miles.utils.dumper_utils import DumperMegatronUtil, DumperPhase
+from miles.utils.types import RolloutBatch
+
+
+def run_forward_backward(
+    args: Namespace,
+    batch_id: int,
+    model: Sequence[DDP],
+    rollout_data: RolloutBatch,
+) -> dict:
+    data_iterator, num_microbatches = get_data_iterator(args, model, rollout_data)
+    assert len(num_microbatches) == 1, "a work unit is a single forward/backward pass"
+
+    for iterator in data_iterator:
+        iterator.reset()
+    for model_chunk in model:
+        model_chunk.train()
+    # disable_optimizer: bf16 loss scaling is the identity, so the pass needs no optimizer
+    setup_train_iteration_config(args, model, None, disable_optimizer=True)
+    reset_grad_metadata_keep_grads(model)
+
+    dumper_phase_util = DumperMegatronUtil(args, model, DumperPhase.FWD_BWD, rollout_id=batch_id)
+    losses_reduced = run_forward_backward_pass(
+        args, dumper_phase_util, data_iterator, model, num_microbatches[0], num_rollouts=None
+    )
+    dumper_phase_util.finalize(model)
+
+    if get_parallel_state().is_pp_last_stage:
+        return aggregate_train_losses(losses_reduced, None)
+    return {}
+
+
+def optim_step(
+    args: Namespace,
+    slot_optimizers: dict[int, SlotOptimizer],
+    adam_params_by_slot: dict[int, dict],
+) -> dict[int, dict]:
+    # batch size 1: grads step as accumulated; normalization is the client's loss weights
+    stepped = {slot: slot_optimizers[slot] for slot in adam_params_by_slot}
+    return step_slot_optimizers(stepped, adam_params_by_slot, clip_grad=args.clip_grad)
+
+
+def load_slot(args: Namespace, model: Sequence[DDP], slot: int, rank: int, alpha: float) -> SlotOptimizer:
+    from megatron.bridge.peft.multi_lora_layers import init_adapter_slot
+
+    init_adapter_slot(model, slot, rank=rank, alpha=alpha)
+    slot_optimizer = SlotOptimizer(args, model, slot)
+    slot_optimizer.zero_grads()
+    return slot_optimizer
+
+
+def unload_slot(model: Sequence[DDP], slot_optimizer: SlotOptimizer) -> None:
+    """The caller drops the instance afterwards; its state dies with it."""
+    from megatron.bridge.peft.multi_lora_layers import clear_adapter_slot
+
+    clear_adapter_slot(model, slot_optimizer.slot)
+    slot_optimizer.zero_grads()