执行摘要

• 一句话：新增 Tinker 每批损失选择、5 种目标函数与 per-datum 输出
• 推荐动作：建议精读。最值得关注的三处设计决策：(1) tinker_losses.py 的「服务端只做 raw sum、归一化交给客户端」模型，以及 loss_normalizer 与执行 batch size 解耦的理由；(2) _sum_loss_and_outputs 在“无监督 token 的 microbatch”上用 logits.sum(dtype=torch.float32) * 0 保持计算图连通的 NaN 防御写法，以及 per-datum 输出 detach 到 CPU 以规避 recompute 重复追加的设计；(3) forward_only 复用同一条 loss pass 而不是另起实现，这是消除双实现漂移的关键取舍。阅读时建议同时对照 review 中“零掩码测试无法验证目标函数”与“per_datum 命名空间污染”两条评论，它们分别解释了最终测试矩阵与 loss.py 守卫的由来；并把 nanjiangwill 关于 O(KN) 全池同步的评论作为独立待办跟踪。

功能与动机

Tinker 协议允许客户端按 work unit 指定目标函数，而 miles 原有实现只有一个由 args.loss_type 决定的全局损失，无法在同一 trainer 上服务不同 loss_fn 的请求。PR body 说明：需要 “Select loss_fn and loss_fn_config per work unit”，并提供 “cross-entropy, importance-sampling, PPO, CISPO, and DRO losses”；归一化由客户端负责，“client loss inputs provide normalization, including custom-loss gradients represented as weights = -dL/dlogprob”；同时要避免全局收集器带来的重复输出——“There is no module-global output collector, so loss recomputation cannot append duplicate outputs”。配套地，“Optimizer requests also apply grad_clip_norm per slot”，并且 “loss_normalizer is separate from the execution batch size, so Tinker accumulation retains raw sums”。本 PR 与 @yushengsu-thu 共同完成，堆叠在 #2840 之上。

实现拆解

1. 新增 Tinker 损失族与注册表（miles/backends/training_utils/loss_hub/tinker_losses.py，新增 163 行）
   - 5 个目标函数：cross_entropy_loss_function、importance_sampling_loss_function、ppo_loss_function、cispo_loss_function、dro_loss_function，统一签名 (args, batch, logits, sum_of_sample_mean) -> (loss, log)。
   - 共享辅助：_target_logprobs 把显式 target_tokens 拼到序列尾部当作 gather label；_as_tensor_like 统一 dtype/device；_response_masks 保证 DP padding datum 的全零掩码不进入目标；_sum_loss_and_outputs 做 raw sum 聚合并产出 per-datum 输出。
   - 默认超参常量 PPO_DEFAULTS / CISPO_DEFAULTS / DRO_DEFAULTS，可被 batch["loss_fn_config"] 逐 batch 覆盖；TINKER_LOSS_FUNCTIONS 是客户端字符串名到实现的注册表。
   - 影响：目标函数从「进程级单例」变成「每 work unit 可换」，这是后续协议层/网关（#2844、#2845）能按请求分派的前提。

2. 每批损失选择与归一化契约（miles/backends/training_utils/loss.py、miles/backends/training_utils/loss_hub/losses.py）
   - get_loss_function(args, loss_fn=None) 在 match args.loss_type 之前短路：只要 loss_fn 非空就返回 TINKER_LOSS_FUNCTIONS[loss_fn]。
   - loss_function 改为 get_loss_function(args, batch.get("loss_fn"))，其余调用点保持兼容。
   - 原来的 global_batch_size = 1（multi-LoRA 特判）被重命名为 loss_normalizer：multi-LoRA 下仍为 1，非 multi-LoRA 回落到 global_batch_size，语义上把“梯度缓冲权重”和“损失归一化分母”两件事拆开。
   - 目标函数返回的 per_datum 被 log.pop("per_datum", None) if batch.get("loss_fn") is not None else None 摘出，并只在 Tinker 路径注入 logging dict，避免污染原生/自定义 loss 的指标命名空间。

3. 训练期数据契约扩展（miles/backends/megatron_utils/model.py）
   - run_forward_backward_pass(..., forward_only=False) 新增参数，前向-only 时用 torch.no_grad() 包裹 forward_backward_func，并透传 forward_only 给 Megatron。
   - forward_step 的 get_batch key 白名单新增 loss_weights、target_tokens、sample_indices，这是 per-datum 损失能读到权重与样本索引的唯一入口。
   - 同一 hunk 顺带把 rollout 前向 forward_only 里的 assert not return_schedule_plan 删掉、并简化了 docstring，属于与本主题无关的约束放宽。
   - 配套：miles/ray/rollout/train_data_conversion.py 的 shard 白名单必须同步补齐，否则 DP 切分后这些向量会静默消失，只会在 trainer 内部才炸——test_delayed_dp_split_preserves_the_tinker_loss_vectors 正是这条守卫。

4. LoRA 执行器与 optim_step 改造（miles/backends/megatron_utils/lora/model.py、lora/actor.py、lora/optimizer.py、miles/ray/train/group.py）
   - run_forward_backward(..., *, forward_only=False)：forward_only 时跳过 reset_grad_metadata_keep_grads，把 per_datum 从各 microbatch 结果里聚合出来，返回 {"metrics": ..., "per_datum": ...}（非 PP 末段返回空 metrics + per_datum）。
   - MultiLoRATrainRayActor.forward_only_logprobs 被 forward_only 取代，直接复用 run_forward_backward(..., forward_only=True)；miles/ray/train/group.py 的 _execute_slots("forward_only", ...) 同步改名，RPC 名字面量变化需要网关侧对齐。
   - optim_step 去掉 args 与 clip_grad，per-slot 的 grad_clip_norm 改由槽位优化器在步进时应用（PR body：“Optimizer requests also apply grad_clip_norm per slot”）。

5. 测试与验证配套
   - 新增 tests/fast/backends/training_utils/loss/test_tinker_loss_outputs.py（142 行）：验证 direct/recomputed 两种模式下 per-datum 输出互不串扰、不重复追加、张量 detach、零掩码 datum 对所有 5 个目标都不产生梯度，并补充 7 组非零目标/梯度用例（IS、PPO 默认与覆盖、CISPO 默认与覆盖、DRO 默认与覆盖）。
   - 新增 tests/fast/backends/training_utils/loss/test_tinker_explicit_targets.py：验证 target 覆盖序列作为 gather label，以及 shifted target 退化为 next-token 打分。
   - tests/fast/backends/megatron_utils/test_step_adapter_slots.py、tests/fast/backends/training_utils/loss/test_training_logprob_reuse.py 随签名调整同步更新。
   - PR body 声明 pre-commit run --all-files 与增量 git diff --check 通过，CPU/GPU CI 尚在排队（“Fresh CPU/GPU CI is pending”），验证与限制记录在 #2846、#3198。

关键文件：

• miles/backends/training_utils/loss_hub/tinker_losses.py（模块 损失函数；类别 source；类型 core-logic；符号 _target_logprobs, _as_tensor_like, _response_masks, _sum_loss_and_outputs）: 本 PR 的核心新增文件：5 个 Tinker 目标函数、共享辅助（_target_logprobs/_response_masks/_sum_loss_and_outputs）与 TINKER_LOSS_FUNCTIONS 注册表，定义了服务端 raw-sum、客户端归一化的损失契约。
• miles/backends/megatron_utils/model.py（模块 训练引擎；类别 source；类型 data-contract；符号 run_forward_backward_pass）: 训练/前向的公共 pipeline 入口：run_forward_backward_pass 新增 forward_only 参数与 no_grad 分支，get_batch 白名单新增 loss_weights/target_tokens/sample_indices，是 per-datum 损失可用的数据契约来源。
• miles/backends/training_utils/loss.py（模块 损失分发；类别 source；类型 core-logic；符号 loss_function）: 损失分发与归一化的中枢：get_loss_function 改为按 batch["loss_fn"] 选择，global_batch_size 抽象为 loss_normalizer，并只在 Tinker 路径摘取 per_datum。
• miles/backends/megatron_utils/lora/model.py（模块 适配器训练；类别 source；类型 data-contract；符号 run_forward_backward, optim_step）: multi-LoRA 执行入口：run_forward_backward 增加 forward_only 关键字参数、跳过梯度元数据重置、聚合 per_datum 并改为返回 {metrics, per_datum}；optim_step 去掉 args/clip_grad。
• miles/backends/megatron_utils/lora/actor.py（模块 适配器训练；类别 source；类型 core-logic；符号 forward_backward, optim_step, forward_only, load_slot）: RPC 面变更：forward_only_logprobs 被 forward_only 取代并复用 run_forward_backward，optim_step 签名与返回值类型调整，直接影响网关调用方。
• tests/fast/backends/training_utils/loss/test_tinker_loss_outputs.py（模块 损失测试；类别 test；类型 test-coverage；符号 test_loss_passes_return_independent_detached_outputs, test_a_zero_loss_mask_removes_the_datum_from_every_objective, test_nonzero_objectives_and_gradients）: 本 PR 的核心测试：覆盖重算不重复输出、detach/隔离、零掩码无梯度，以及 7 组非零目标与梯度用例（对应 reviewer 要求补齐的验证）。
• miles/backends/training_utils/loss_hub/losses.py（模块 损失注册；类别 source；类型 core-logic；符号 get_loss_function）: get_loss_function 增加 loss_fn 参数并在 match args.loss_type 之前短路到 TINKER_LOSS_FUNCTIONS，是“每批损失选择”的接入口。
• tests/fast/backends/training_utils/loss/test_tinker_explicit_targets.py（模块 目标对齐测试；类别 test；类型 test-coverage；符号 _naive_compute_log_probs, _target_logprobs, test_targets_replace_the_sequence_as_gather_labels, test_shifted_targets_degenerate_to_next_token_scoring）: 锁定 target_tokens 作为显式 gather label 的语义（含 shifted target 退化为 next-token 打分），是 Tinker 客户端契约的回归护栏。
• miles/ray/rollout/train_data_conversion.py（模块 数据切分；类别 source；类型 data-contract）: DP 切分的 shard key 白名单需与 get_batch 同步；漏写会让 loss_weights/advantages/rollout_log_probs 静默消失，只在 trainer 内部报错。
• miles/backends/megatron_utils/lora/optimizer.py（模块 适配器优化；类别 source；类型 configuration）: 配合 optim_step 瘦身，把 grad_clip_norm 改为按 slot 应用，是归一化职责从训练端迁移到优化器步骤的一部分。

关键符号：cross_entropy_loss_function, importance_sampling_loss_function, ppo_loss_function, cispo_loss_function, dro_loss_function, _target_logprobs, _response_masks, _sum_loss_and_outputs, get_loss_function, loss_function, run_forward_backward_pass, run_forward_backward, optim_step, forward_only

关键源码片段

miles/backends/training_utils/loss_hub/tinker_losses.py

本 PR 的核心新增文件：5 个 Tinker 目标函数、共享辅助（_target_logprobs/_response_masks/_sum_loss_and_outputs）与 TINKER_LOSS_FUNCTIONS 注册表，定义了服务端 raw-sum、客户端归一化的损失契约。

# miles/backends/training_utils/loss_hub/tinker_losses.py
# Tinker 损失族：每个 work unit 由客户端指定 loss_fn，训练端只累加 raw sum，
# 归一化由客户端在 loss_weights 里给出（自定义损失的梯度表示为 weights = -dL/dlogprob）。
​
PPO_DEFAULTS = {"clip_low_threshold": 0.8, "clip_high_threshold": 1.2}
CISPO_DEFAULTS = {"clip_low_threshold": 0.0, "clip_high_threshold": 4.0}
DRO_DEFAULTS = {"beta": 0.05}
​
​
def _target_logprobs(args: Namespace, batch: RolloutBatch, logits: torch.Tensor) -> list[torch.Tensor]:
    # Tinker 的 target 是显式标签：用它替换序列 response 区段的 token，
    # 于是 gather 的 label 不再等于输入 token 本身（rl_loop 风格）。
    label_tokens = [
        torch.cat([tokens[: len(tokens) - len(targets)], _as_tensor_like(targets, tokens)])
        for tokens, targets in zip(batch["unconcat_tokens"], batch["target_tokens"], strict=True)
    ]
    outputs = get_log_probs_and_entropy(
        logits,
        args=args,
        unconcat_tokens=label_tokens,
        total_lengths=batch["total_lengths"],
        response_lengths=batch["response_lengths"],
        with_entropy=False,
        max_seq_lens=batch.get("max_seq_lens", None),
    )
    return outputs["log_probs"]
​
​
def _sum_loss_and_outputs(batch, logits, log_probs, per_datum_losses):
    if per_datum_losses:
        # 逐 datum 求和后直接相加：服务端不做除以 batch size 的归一化
        loss = torch.stack(per_datum_losses).sum()
    else:
        # 没有监督 token 的 microbatch 也要让计算图保持连通；
        # fp32 求和避免 fp16 溢出成 inf，0 * inf 会变成 nan
        loss = logits.sum(dtype=torch.float32) * 0
    per_datum = [
        # 输出一律 detach 并搬到 CPU：它只用于回传给客户端，不参与反传
        {"sample_index": index, "logprobs": log_prob.detach().cpu(), "loss": sample_loss.detach().cpu()}
        for index, log_prob, sample_loss in zip(batch["sample_indices"], log_probs, per_datum_losses, strict=True)
    ]
    return loss, {"loss": loss.detach(), "per_datum": per_datum}
​
​
def ppo_loss_function(args, batch, logits, sum_of_sample_mean):
    # 裁剪阈值允许被客户端的 loss_fn_config 逐 batch 覆盖
    config = batch.get("loss_fn_config") or {}
    clip_low = config.get("clip_low_threshold", PPO_DEFAULTS["clip_low_threshold"])
    clip_high = config.get("clip_high_threshold", PPO_DEFAULTS["clip_high_threshold"])
    log_probs = _target_logprobs(args, batch, logits)
    per_datum_losses = []
    for log_prob, sampling_log_prob, advantage, mask in zip(
        log_probs, batch["rollout_log_probs"], batch["advantages"], _response_masks(batch, log_probs), strict=True
    ):
        ratio = torch.exp(log_prob - _as_tensor_like(sampling_log_prob, log_prob))
        advantages = _as_tensor_like(advantage, log_prob)
        objective = torch.minimum(ratio * advantages, torch.clamp(ratio, clip_low, clip_high) * advantages)
        per_datum_losses.append(-(objective * mask).sum())
    return _sum_loss_and_outputs(batch, logits, log_probs, per_datum_losses)
​
​
# 客户端 loss_fn 字符串 -> 实现；loss.py 通过 batch["loss_fn"] 查这张表
TINKER_LOSS_FUNCTIONS = {
    "cross_entropy": cross_entropy_loss_function,
    "importance_sampling": importance_sampling_loss_function,
    "ppo": ppo_loss_function,
    "cispo": cispo_loss_function,
    "dro": dro_loss_function,
}

miles/backends/training_utils/loss.py

损失分发与归一化的中枢：get_loss_function 改为按 batch["loss_fn"] 选择，global_batch_size 抽象为 loss_normalizer，并只在 Tinker 路径摘取 per_datum。

# miles/backends/training_utils/loss.py —— loss_function 的核心分支
# 损失函数改为逐 batch 选择：同一个 trainer 可服务不同 loss_fn 的 work unit
func = get_loss_function(args, batch.get("loss_fn"))
​
if args.recompute_loss_function:
    # checkpoint 重算时只重跑损失，不会重复产出 per-datum 结果
    loss, log = checkpoint(func, args, batch, logits, sum_of_sample_mean)
else:
    loss, log = func(args, batch, logits, sum_of_sample_mean)
​
# 保持各 rank 的计算图连通，避免 CP/allgather 下反向不一致导致挂起
if parallel_state.cp.size > 1 and args.allgather_cp:
    loss = loss + 0 * logits.sum(dtype=torch.float32)
​
if num_rollouts is not None:
    global_batch_size = num_rollouts
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)
# Tinker 损失已带客户端归一化；multi-LoRA 下样本以权重 1 进入梯度缓冲，
# per-adapter 归一化推迟到 optimizer step，所以这里不能再除一次
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 / loss_normalizer * loss_parallel_size
    else:
        loss = loss / loss_normalizer * parallel_state.intra_dp.size
else:
    if apply_megatron_loss_scaling:
        loss = loss * parallel_state.cp.size
​
# per_datum 只属于 Tinker 协议；原生损失与 --custom-loss-function-path
# 的指标命名空间不预留这个键，因此只在 Tinker 路径上摘取
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()),
            device=logits.device,
        ),
    },
)

miles/backends/megatron_utils/lora/model.py

multi-LoRA 执行入口：run_forward_backward 增加 forward_only 关键字参数、跳过梯度元数据重置、聚合 per_datum 并改为返回 {metrics, per_datum}；optim_step 去掉 args/clip_grad。

# miles/backends/megatron_utils/lora/model.py
# Tinker 的 forward() 与 forward_backward() 共用同一条 loss pass
def run_forward_backward(
    args,
    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"
    # 模型切 chunk、train() 模式切换等准备步骤与训练态一致
    setup_train_iteration_config(args, model, None, disable_optimizer=True)
    if not forward_only:
        # forward() 不产生梯度，因此不必保留上一轮的梯度元数据
        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,
        forward_only=forward_only,
    )
    # per-datum 输出随 microbatch 结果一起回流，不再有模块级全局收集器，
    # 因此 recompute_loss_function 重算损失时不会把同一 datum 追加两次
    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 {"metrics": aggregate_train_losses(losses_reduced, None), "per_datum": per_datum_outputs}
    return {"metrics": {}, "per_datum": per_datum_outputs}
​
​
def optim_step(slot_optimizers, adam_params_by_slot):
    # batch size 为 1：梯度按累积值直接步进，归一化由客户端损失权重承担，
    # 每个 slot 的 grad_clip_norm 在槽位优化器内部应用
    stepped = {slot: slot_optimizers[slot] for slot in adam_params_by_slot}
    return step_slot_optimizers(stepped, adam_params_by_slot)

评论区精华

本 PR 共 5 条 review comment、7 条 review，实际技术交锋来自 nanjiangwill 的 4 条评论，yushengsu-thu 最终 APPROVED。

1) 测试有效性（已被采纳修复）。nanjiangwill 在 test_tinker_loss_outputs.py:55 指出早期的参数化用例“cannot validate the four new objectives because every result is multiplied by an all-zero mask; even an implementation that always returns zero passes”，要求补齐 IS/PPO/CISPO/DRO 的非零 loss 与梯度、包含默认与覆盖的 clip/beta。最终版本新增了 test_nonzero_objectives_and_gradients 的 7 组参数化用例，给出逐 token 期望值与梯度断言，问题闭环。

2) 指标命名空间污染（已被采纳修复）。针对 loss.py 的 per_datum = log.pop("per_datum", None)，他指出该分支“also runs for native and custom miles losses, where batch["loss_fn"] is None … an existing metric with this name now disappears from logging”，建议只在 Tinker 路径摘取。最终实现改为 if batch.get("loss_fn") is not None else None，并在 run_forward_backward 里用 log.pop("per_datum", None) 做双重防护。

3) per-datum schema 的可扩展性（非阻塞、未落代码）。在 tinker_losses.py:59 他提醒：“future client-side losses may also request multi-target logprobs, values, or entropy … just avoid treating {sample_index, logprobs, loss} as a closed schema”。结论是当前实现不改，但后续客户端损失上线时要预留字段扩展。

4) 优化器步进性能（本 PR 未解决）。在 miles/backends/megatron_utils/lora/executor.py 上他指出：“With DP > 1, this loop turns one optimizer barrier into K full-pool synchronizations … Stepping K of N slots therefore moves O(KN) adapter parameters—O(N²) when K=N”，建议把 per-slot clip 值合并进一次 step_adapter_slots 调用，只做一次 outcome merge 与参数 all-gather。该反馈对应到本 PR 的 optim_step 瘦身（去掉 args/clip_grad、per-slot grad_clip_norm），但 O(KN) 的同步模式在本 PR 的改动里看不到收敛。

5) yushengsu-thu 在 tinker_losses.py 上留了一条 “Suggestion: https://github.com/radixark/miles/pull/3145” 的跟进指向，未在本 PR 展开。

• 零掩码参数化测试无法验证 IS/PPO/CISPO/DRO 目标函数 (testing): 已采纳：新增 test_nonzero_objectives_and_gradients，7 组参数化用例给出逐 token 期望值与梯度断言（is / ppo-default / ppo-override / cispo-default / cispo-override / dro-default / dro-override）。
• per_datum 摘取会删除原生/自定义损失的同名指标 (correctness): 已采纳：改为 per_datum = log.pop("per_datum", None) if batch.get("loss_fn") is not None else None，并在 run_forward_backward 中再次用 log.pop 兜底。
• per-datum 输出 schema 的可扩展性 (design): 非阻塞：当前仅保留 sample_index/logprobs/loss，后续客户端损失上线时需按该建议扩展字段而不破坏现有消费方。
• DP > 1 时逐 slot 步进造成 O(KN) 参数搬运与全池同步 (performance): 本 PR 未收敛：optim_step 已去掉 args/clip_grad 并把裁剪下沉到 slot 优化器，但逐 slot 同步模式仍存在，需要后续按该建议合并调用。
• 跟进建议指向 PR#3145 (other): 作为后续跟进记录，本 PR 内无对应代码改动；上下文未提供该 PR 标题与内容，无法判断具体后续范围。

风险与影响

• 风险：
  • 归一化语义变化（回归风险，最高）：loss.py 把 global_batch_size 替换为 loss_normalizer，multi-LoRA 下恒为 1，即训练端只累加客户端给的 raw sum。如果客户端对某个 loss_fn 没有自带归一化（例如直接传 advantage、或自定义损失只给 weights = -dL/dlogprob 而没缩放），梯度尺度会整体偏移但不会报错，属于静默失配；PR body 也承认 CI 尚待执行，建议在 GPU CI 上做一次 loss 尺度对照（与单卡小批量手算比对）。
  • per_datum 的主机内存与序列化开销：_sum_loss_and_outputs 对每个 datum 都执行 log_prob.detach().cpu() 与 sample_loss.detach().cpu()，且在 lora/model.py 中是在 is_pp_last_stage 判断之前就做全量列表推导，非末段同样构造（随后被丢弃）。长序列/大 batch 下这部分 CPU 拷贝与后续 RPC 序列化会成为新开销，需要关注。
  • PP 非末段路径缺乏测试覆盖：run_forward_backward 返回 {"metrics": {}, "per_datum": per_datum_outputs} 的分支里，per_datum_outputs 直接迭代 losses_reduced 取 microbatch["per_datum"]；非末段的 losses_reduced 结构与末段并不相同，而现有 fast 测试都是单 stage，建议在 TP/PP>1 的 GPU CI 或 e2e 上确认该分支不会因结构差异报错。
  • target_tokens 长度契约无断言：_target_logprobs 用 tokens[: len(tokens) - len(targets)] 拼接，隐含 len(targets) == response_lengths；长度不匹配时不会抛错，只会让 gather 位置整体偏移，建议补断言或校验。
  • 顺带删除的断言：miles/backends/megatron_utils/model.py 中 rollout 前向 forward_only 的 assert not return_schedule_plan 被移除，属于与本 PR 主题无关的约束放宽，review 中未见讨论，建议确认是否有意为之。
  • 优化器步进性能隐患：nanjiangwill 指出的 K 次 step_adapter_slots 触发 O(KN) 参数 all-gather / 全池同步问题在本 PR 未闭环，在 DP > 1 且多 slot 并发步进时会放大。
  • 耦合与兼容性：actor RPC 名称 forward_only_logprobs -> forward_only 是字面量级的接口变更，任何仍按旧名调用的网关/脚本都会直接失败；optim_step 去掉 args 后，per-slot clip 语义依赖优化器内部实现，需要与其他 multi-LoRA 分片（3/7 的槽位检查点、优化器状态）保持同步。
• 影响：
  • 对用户：Tinker 客户端现在可以按 work unit 选择 5 种目标函数并通过 forward() 拿到逐 datum 的 sample_index / logprobs / loss，为 rl_loop 风格的显式 target、自定义损失（weights = -dL/dlogprob）打开通路；原生 miles RL/SFT 用户在 batch["loss_fn"] 为 None 时走旧路径，理论上无行为变化（per_datum 摘取有守卫），但 multi-LoRA 的归一化契约变化对所有使用该链路的实验都生效。
  • 对系统：loss_function 的返回契约从 {keys, values} 扩展出可选的 per_datum；Megatron 训练循环新增 forward_only 参数与 torch.no_grad() 分支；LoRA 执行器新增一条与 forward_backward 共用 loss pass 的推理路径，减少了此前 compute_log_prob 与训练 loss 两条实现漂移的可能。
  • 对团队/后续演进：本 PR 是 7 篇系列第 2 篇，3/7~7/7（槽位检查点、快照发布、协议层、网关 runtime、示例与测试）都依赖这里的 loss_fn 分派、loss_normalizer、per_datum 结构与 actor 方法名；review 中暴露的优化器 all-gather 复杂度问题会成为后续优化器步骤统一改造的输入。
  • 风险标记：核心训练路径变更, 归一化契约调整, 新增 no_grad 前向路径, PP 多段路径缺测试, per_datum 主机内存开销, 优化器逐 slot 步进性能隐患, CI 尚未验证

关联脉络

• PR #2845 [multi-lora v2 6/7] Tinker gateway runtime and entry point: 同一系列下一篇，网关 runtime 依赖本 PR 定下的 per_datum 输出结构与 forward_only/optim_step 的 actor 接口。
• PR #2844 [multi-lora v2 5/7] Tinker gateway protocol face: 协议层按 work unit 组装 loss_fn / loss_fn_config 与 sample_index，正是本 PR 损失选择契约的消费方。
• PR #2842 [multi-lora v2 3/7] Per-slot checkpoints with optimizer state: 同为 slot 优化器线路：本 PR 调整 optim_step 与 per-slot grad_clip_norm，与本系列检查点/优化器状态改动直接相邻。
• PR #3245 fix(megatron): respect --no-load-optim when resuming a LoRA adapter: 同样涉及 LoRA 优化器状态与恢复语义，与本 PR 的 optim_step 归一化/裁剪职责调整属于同一模块的演进。
• PR #3083 fix(megatron): respect --no-save-optim when saving checkpoints: LoRA 检查点与优化器状态保存的历史修复，说明该线路对优化器状态语义变更敏感，本 PR 的 grad_clip_norm 下沉需与其保持一致。