[multi-lora v2 2/7] Per-batch loss selection and Tinker losses

miles/backends/megatron_utils/lora/actor.py

@@ -1,12 +1,9 @@
 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.backends.training_utils.data import get_rollout_data
 from miles.utils.ray_utils import Box
 from miles.utils.tracking_utils.structured_log import with_logs
 
@@ -24,21 +21,19 @@ def forward_backward(self, batch_id: int, rollout_data_ref: Box) -> dict:
             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]:
+    def optim_step(self, adam_params_by_slot: dict[int, dict]) -> dict[int, dict]:
         self._heartbeat.bump()
-        return lora_model.optim_step(self.args, self.slot_optimizers, adam_params_by_slot)
+        return lora_model.optim_step(self.slot_optimizers, adam_params_by_slot)
 
     @with_logs
-    def forward_only_logprobs(self, batch_id: int, rollout_data_ref: Box) -> Box | None:
+    def forward_only(self, batch_id: int, rollout_data_ref: Box) -> dict:
+        """Same loss pass as forward_backward, without the backward: the Tinker
+        forward() contract returns the requested loss per datum."""
         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()}))
+            return lora_model.run_forward_backward(self.args, batch_id, self.model, rollout_data, forward_only=True)
 
     @with_logs
     def load_slot(self, slot: int, rank: int, alpha: float) -> None:

miles/backends/megatron_utils/lora/model.py

@@ -23,6 +23,8 @@ def run_forward_backward(
     batch_id: int,
     model: Sequence[DDP],
     rollout_data: RolloutBatch,
+    *,
+    forward_only: bool = False,
 ) -> 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"
@@ -33,27 +35,33 @@ def run_forward_backward(
         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)
+    if not forward_only:
+        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
+        args,
+        dumper_phase_util,
+        data_iterator,
+        model,
+        num_microbatches[0],
+        num_rollouts=None,
+        forward_only=forward_only,
     )
+    per_datum_outputs = [output for microbatch in losses_reduced for output in microbatch["per_datum"]]
     dumper_phase_util.finalize(model)
 
     if get_parallel_state().is_pp_last_stage:
-        return aggregate_train_losses(losses_reduced, None)
-    return {}
+        return {"metrics": aggregate_train_losses(losses_reduced, None), "per_datum": per_datum_outputs}
+    return {"metrics": {}, "per_datum": per_datum_outputs}
 
 
 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)
+    return step_slot_optimizers(stepped, adam_params_by_slot)
 
 
 def load_slot(args: Namespace, model: Sequence[DDP], slot: int, rank: int, alpha: float) -> SlotOptimizer:

miles/backends/megatron_utils/lora/optimizer.py

@@ -192,16 +192,16 @@ def reset_grad_metadata_keep_grads(model_chunks) -> None:
 def step_slot_optimizers(
     slot_optimizers: dict[int, SlotOptimizer],
     adam_params_by_slot: dict[int, dict],
-    *,
-    clip_grad: float,
 ) -> dict[int, dict]:
     """Step slots in the same order on every rank; execution failures invalidate the cell."""
     slots = sorted(adam_params_by_slot)
     for slot in slots:
         slot_optimizers[slot].apply_adam_params(adam_params_by_slot[slot])
         slot_optimizers[slot].prepare_grads()
 
-    outcomes = {slot: slot_optimizers[slot].clip_and_step(clip_grad) for slot in slots}
+    outcomes = {
+        slot: slot_optimizers[slot].clip_and_step(adam_params_by_slot[slot]["grad_clip_norm"]) for slot in slots
+    }
     for slot in slots:
         slot_optimizers[slot].zero_grads()
     for slot in slots:

miles/backends/megatron_utils/model.py

@@ -302,21 +302,8 @@ def forward_only(
     def forward_step(
         data_iterator: DataIterator, model: GPTModel, return_schedule_plan: bool = False
     ) -> tuple[torch.Tensor, Callable[[torch.Tensor], dict[str, list[torch.Tensor]]]]:
-        """Forward step used by Megatron's pipeline engine.
+        """Return the model output and its batch-bound loss callback."""
 
-        Args:
-            data_iterator (DataIterator): Input data iterator.
-            model (GPTModel): The GPT model chunk to execute.
-
-        Returns:
-            tuple[torch.Tensor, Callable[[torch.Tensor], dict[str, list[torch.Tensor]]]]:
-            Output tensor(s) and a callable that computes and packages results
-            to be collected by the engine.
-        """
-
-        assert not return_schedule_plan, "forward_only step should never return schedule plan"
-
-        # Get the batch.
         batch = get_batch(
             data_iterator,
             [
@@ -415,13 +402,15 @@ def _zero_grads(model: Sequence[DDP], optimizer: MegatronOptimizer | None, disab
         optimizer.zero_grad()
 
 
-def run_forward_backward_pass(args, dumper_phase_util, data_iterator, model, num_microbatches, num_rollouts):
+def run_forward_backward_pass(
+    args, dumper_phase_util, data_iterator, model, num_microbatches, num_rollouts, forward_only=False
+):
     """One pipeline forward/backward pass over the microbatches; no optimizer interaction."""
 
     @dumper_phase_util.wrap_forward_step
     def forward_step(data_iterator: DataIterator, model: GPTModel, return_schedule_plan: bool = False) -> tuple[
         torch.Tensor,
-        Callable[[torch.Tensor], tuple[torch.Tensor, int, dict[str, torch.Tensor | list[str]]]],
+        Callable[[torch.Tensor], tuple[torch.Tensor, int, dict]],
     ]:
         """Forward step used by Megatron's pipeline engine during training.
 
@@ -430,7 +419,7 @@ def forward_step(data_iterator: DataIterator, model: GPTModel, return_schedule_p
             model (GPTModel): The GPT model chunk to execute.
 
         Returns:
-            tuple[torch.Tensor, Callable[[torch.Tensor], tuple[torch.Tensor, int, dict[str, torch.Tensor | list[str]]]]]:
+            tuple[torch.Tensor, Callable[[torch.Tensor], tuple[torch.Tensor, int, dict]]]:
             Output tensor(s) and the loss function, which returns
             (loss, num_elems, {"keys": list[str], "values": torch.Tensor}).
         """
@@ -455,6 +444,9 @@ def forward_step(data_iterator: DataIterator, model: GPTModel, return_schedule_p
                 "witness_ids",
                 "opd_reverse_kl",
                 "rollout_mask_sums",
+                "loss_weights",
+                "target_tokens",
+                "sample_indices",
             ],
             args.data_pad_size_multiplier,
             args.qkv_format,
@@ -515,16 +507,17 @@ def forward_step(data_iterator: DataIterator, model: GPTModel, return_schedule_p
 
     # Forward pass.
     forward_backward_func = get_forward_backward_func()
-    return forward_backward_func(
-        forward_step_func=forward_step,
-        data_iterator=data_iterator,
-        model=model,
-        num_microbatches=num_microbatches,
-        seq_length=args.seq_length,
-        micro_batch_size=args.micro_batch_size,
-        decoder_seq_length=args.decoder_seq_length,
-        forward_only=False,
-    )
+    with torch.no_grad() if forward_only else nullcontext():
+        return forward_backward_func(
+            forward_step_func=forward_step,
+            data_iterator=data_iterator,
+            model=model,
+            num_microbatches=num_microbatches,
+            seq_length=args.seq_length,
+            micro_batch_size=args.micro_batch_size,
+            decoder_seq_length=args.decoder_seq_length,
+            forward_only=forward_only,
+        )
 
 
 def train_one_step(

miles/backends/training_utils/data.py

@@ -210,6 +210,10 @@ def get_batch(
     if "dynamic_global_batch_size" in data_iterator.rollout_data:
         batch["dynamic_global_batch_size"] = data_iterator.rollout_data["dynamic_global_batch_size"]
 
+    for unit_key in ("loss_fn", "loss_fn_config"):
+        if unit_key in data_iterator.rollout_data:
+            batch[unit_key] = data_iterator.rollout_data[unit_key]
+
     # No-op safety net if batches reach get_batch without rollout-level preprocessing.
     expand_multimodal_rollout_data_in_place(batch, qkv_format=qkv_format)
 

miles/backends/training_utils/loss.py

@@ -132,11 +132,12 @@ def loss_function(
     logits: torch.Tensor,
     apply_megatron_loss_scaling: bool = False,
     num_rollouts: int | None = None,
-) -> tuple[torch.Tensor, int | torch.Tensor, dict[str, list[str] | torch.Tensor]]:
+) -> tuple[torch.Tensor, int | torch.Tensor, dict]:
     """Dispatch to the configured loss and rescale for Megatron integration.
 
     Selects one of "policy_loss", "value_loss", "sft_loss", or a custom loss
-    function based on `args.loss_type`, computes the loss and metrics, then
+    function based on `args.loss_type`, or the Tinker loss specified by
+    `batch["loss_fn"]`. Computes the loss and metrics, then
     rescales the loss by micro-batch and parallelism factors to integrate with
     Megatron's gradient accumulation.
 
@@ -157,6 +158,7 @@ def loss_function(
           `args.calculate_per_token_loss` is True, else `1` (int).
         - `logging_dict` has keys "keys" (list of str metric names) and
           "values" (1D tensor: [count, metric1, metric2, ...]).
+          Tinker losses may also return "per_datum" outputs.
     """
     parallel_state = get_parallel_state()
     num_tokens = sum([torch.clamp_min(loss_mask.sum(), 1) for loss_mask in batch["loss_masks"]])
@@ -172,7 +174,7 @@ def loss_function(
         denominators=batch.get("rollout_mask_sums", None),
     )
 
-    func = get_loss_function(args)
+    func = get_loss_function(args, batch.get("loss_fn"))
 
     if args.recompute_loss_function:
         loss, log = checkpoint(
@@ -196,29 +198,28 @@ def loss_function(
     else:
         assert args.use_dynamic_global_batch_size == ("dynamic_global_batch_size" in batch)
         global_batch_size = batch.get("dynamic_global_batch_size", args.global_batch_size)
-    # Multi-LoRA: samples enter the gradient buffers with weight 1; per-adapter
-    # normalization (1/adapter_global_batch_size, a constant known in advance)
-    # is applied to the accumulated slot gradient at optimizer-step time.
-    if is_multi_lora_enabled(args):
-        global_batch_size = 1
+    # Tinker losses already carry the client's normalization
+    loss_normalizer = 1 if is_multi_lora_enabled(args) else global_batch_size
     if not args.calculate_per_token_loss:
         if apply_megatron_loss_scaling:
             loss_parallel_size = (
                 parallel_state.intra_dp.size
                 if args.true_on_policy_mode and parallel_state.is_ulysses_cp
                 else parallel_state.intra_dp_cp.size
             )
-            loss = loss * num_microbatches / global_batch_size * loss_parallel_size
+            loss = loss * num_microbatches / loss_normalizer * loss_parallel_size
         else:
-            loss = loss / global_batch_size * parallel_state.intra_dp.size
+            loss = loss / loss_normalizer * parallel_state.intra_dp.size
     else:
         if apply_megatron_loss_scaling:
             loss = loss * parallel_state.cp.size
 
+    per_datum = log.pop("per_datum", None) if batch.get("loss_fn") is not None else None
     return (
         loss,
         torch.tensor(num_tokens if args.calculate_per_token_loss else 1, device=logits.device),
         {
+            **({"per_datum": per_datum} if per_datum is not None else {}),
             "keys": list(log.keys()),
             "values": torch.tensor(
                 [num_samples if not args.calculate_per_token_loss else num_tokens] + list(log.values()),

miles/backends/training_utils/loss_hub/losses.py

@@ -505,7 +505,11 @@ def sft_loss_function(
     )
 
 
-def get_loss_function(args: Namespace) -> LossFunction:
+def get_loss_function(args: Namespace, loss_fn: str | None = None) -> LossFunction:
+    if loss_fn is not None:
+        from miles.backends.training_utils.loss_hub.tinker_losses import TINKER_LOSS_FUNCTIONS
+
+        return TINKER_LOSS_FUNCTIONS[loss_fn]
     match args.loss_type:
         case "policy_loss":
             return policy_loss_function