769 lines
30 KiB
Python
769 lines
30 KiB
Python
# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
|
|
# SPDX-License-Identifier: Apache-2.0
|
|
|
|
import logging
|
|
from argparse import Namespace
|
|
from collections.abc import Callable, Iterator
|
|
from typing import Any
|
|
|
|
import torch
|
|
|
|
logger = logging.getLogger(__name__)
|
|
from megatron.core import mpu
|
|
from torch.utils.checkpoint import checkpoint
|
|
|
|
from slime.utils.distributed_utils import distributed_masked_whiten
|
|
from slime.utils.misc import load_function
|
|
from slime.utils.ppo_utils import (
|
|
calculate_log_probs_and_entropy,
|
|
compute_approx_kl,
|
|
compute_gspo_kl,
|
|
compute_opsm_mask,
|
|
compute_policy_loss,
|
|
get_advantages_and_returns_batch,
|
|
get_grpo_returns,
|
|
get_reinforce_plus_plus_baseline_advantages,
|
|
get_reinforce_plus_plus_returns,
|
|
)
|
|
from slime.utils.types import RolloutBatch
|
|
|
|
from .cp_utils import all_gather_with_cp, get_logits_and_tokens_offset_with_cp, get_sum_of_sample_mean
|
|
|
|
|
|
def get_responses(
|
|
logits: torch.Tensor,
|
|
*,
|
|
args: Namespace,
|
|
unconcat_tokens: list[torch.Tensor],
|
|
total_lengths: list[int],
|
|
response_lengths: list[int],
|
|
) -> Iterator[tuple[torch.Tensor, torch.Tensor]]:
|
|
"""Yield response-aligned `(logits_chunk, tokens_chunk)` pairs per sample.
|
|
|
|
After squeezing batch dimension and applying temperature scaling, this
|
|
function extracts the logits and tokens corresponding to response segments
|
|
for each sample. When context parallelism is disabled, it slices directly
|
|
from the concatenated sequence. With context parallelism enabled, it
|
|
handles split sequences across ranks.
|
|
|
|
Args:
|
|
logits: Model outputs with shape `[1, T, V]` (policy) or `[1, T, 1]`
|
|
(value). Must be float32.
|
|
args: Configuration containing `rollout_temperature` for scaling.
|
|
unconcat_tokens: List of token tensors (prompt+response) per sample.
|
|
total_lengths: Total sequence lengths (prompt+response) per sample.
|
|
response_lengths: Response segment lengths per sample.
|
|
|
|
Yields:
|
|
Tuple of `(logits_chunk, tokens_chunk)` where `logits_chunk` is shape
|
|
`[R, V]` (policy) or `[R, 1]` (value) and `tokens_chunk` is shape `[R]`
|
|
(1D int64), both aligned to response tokens for one sample.
|
|
"""
|
|
assert logits.size(0) == 1, f"{logits.shape}"
|
|
assert logits.dtype == torch.float32, f"{logits.dtype}"
|
|
|
|
logits = logits.squeeze(0)
|
|
logits = logits.div(args.rollout_temperature)
|
|
|
|
cp_size = mpu.get_context_parallel_world_size()
|
|
end = 0
|
|
for tokens, total_length, response_length in zip(unconcat_tokens, total_lengths, response_lengths, strict=False):
|
|
if cp_size == 1:
|
|
end += total_length
|
|
start = end - response_length
|
|
logits_chunk = logits[start - 1 : end - 1]
|
|
tokens_chunk = tokens[-response_length:]
|
|
else:
|
|
# TODO: this is super ugly... do better abstraction.
|
|
chunk_size, chunks_offset, logits_offset, tokens_offset = get_logits_and_tokens_offset_with_cp(
|
|
total_length, response_length
|
|
)
|
|
|
|
logits_0, logits_1 = logits[end : end + chunk_size], logits[end + chunk_size : end + 2 * chunk_size]
|
|
end += 2 * chunk_size
|
|
|
|
logits_0 = logits_0[logits_offset[0][0] - chunks_offset[0][0] : logits_offset[0][1] - chunks_offset[0][0]]
|
|
tokens_0 = tokens[tokens_offset[0][0] : tokens_offset[0][1]]
|
|
|
|
logits_1 = logits_1[logits_offset[1][0] - chunks_offset[1][0] : logits_offset[1][1] - chunks_offset[1][0]]
|
|
tokens_1 = tokens[tokens_offset[1][0] : tokens_offset[1][1]]
|
|
|
|
assert logits_0.size(0) == tokens_0.size(0), f"{logits_0.size(0)} vs {tokens_0.size(0)}"
|
|
assert logits_1.size(0) == tokens_1.size(0), f"{logits_1.size(0)} vs {tokens_1.size(0)}"
|
|
|
|
logits_chunk = torch.cat([logits_0, logits_1], dim=0)
|
|
tokens_chunk = torch.cat([tokens_0, tokens_1], dim=0)
|
|
|
|
yield logits_chunk, tokens_chunk
|
|
|
|
|
|
def get_log_probs_and_entropy(
|
|
logits: torch.Tensor,
|
|
*,
|
|
args: Namespace,
|
|
unconcat_tokens: list[torch.Tensor],
|
|
total_lengths: list[int],
|
|
response_lengths: list[int],
|
|
with_entropy: bool = False,
|
|
non_loss_data: bool = True,
|
|
) -> dict[str, list[torch.Tensor]]:
|
|
"""Compute per-token log-probabilities (and optionally entropy) on responses.
|
|
|
|
For each sample, extracts response-aligned logits and tokens, then computes
|
|
log-probabilities via softmax across the tensor-parallel group. Log-probs
|
|
are squeezed from `[R, 1]` to `[R]`. Entropy values are always appended
|
|
(even when `with_entropy=False`), but only included in the result dict
|
|
when requested.
|
|
|
|
Args:
|
|
logits: Policy logits with shape `[1, T, V]`.
|
|
args: Configuration (temperature applied in `get_responses`).
|
|
unconcat_tokens: List of token tensors per sample.
|
|
total_lengths: Total sequence lengths per sample.
|
|
response_lengths: Response segment lengths per sample.
|
|
with_entropy: If True, include "entropy" key in result.
|
|
non_loss_data: Unused; kept for API compatibility.
|
|
|
|
Returns:
|
|
Dict with key "log_probs" mapping to a list of `[R]` tensors per
|
|
sample. If `with_entropy` is True, also includes "entropy" key with
|
|
a list of `[R]` tensors.
|
|
"""
|
|
assert non_loss_data
|
|
log_probs_list = []
|
|
entropy_list = []
|
|
for logits_chunk, tokens_chunk in get_responses(
|
|
logits,
|
|
args=args,
|
|
unconcat_tokens=unconcat_tokens,
|
|
total_lengths=total_lengths,
|
|
response_lengths=response_lengths,
|
|
):
|
|
log_prob, entropy = calculate_log_probs_and_entropy(
|
|
logits_chunk,
|
|
tokens_chunk,
|
|
mpu.get_tensor_model_parallel_group(),
|
|
with_entropy=with_entropy,
|
|
chunk_size=args.log_probs_chunk_size,
|
|
)
|
|
|
|
log_probs_list.append(log_prob.squeeze(-1))
|
|
entropy_list.append(entropy)
|
|
|
|
res = {
|
|
"log_probs": log_probs_list,
|
|
}
|
|
if with_entropy:
|
|
res["entropy"] = entropy_list
|
|
return res
|
|
|
|
|
|
def get_values(
|
|
logits: torch.Tensor,
|
|
*,
|
|
args: Namespace,
|
|
unconcat_tokens: list[torch.Tensor],
|
|
total_lengths: list[int],
|
|
response_lengths: list[int],
|
|
with_entropy: bool = False,
|
|
non_loss_data: bool = True,
|
|
) -> dict[str, list[torch.Tensor]]:
|
|
"""Extract per-token value predictions over response tokens.
|
|
|
|
For each sample, extracts response-aligned chunks from the value head
|
|
output and squeezes the final dimension from `[R, 1]` to `[R]`.
|
|
|
|
Args:
|
|
logits: Value head output with shape `[1, T, 1]`.
|
|
args: Configuration (passed to `get_responses` which uses
|
|
`rollout_temperature` even though values don't need temperature).
|
|
unconcat_tokens: List of token tensors per sample.
|
|
total_lengths: Total sequence lengths per sample.
|
|
response_lengths: Response segment lengths per sample.
|
|
with_entropy: Unused; kept for signature compatibility.
|
|
non_loss_data: Unused; kept for signature compatibility.
|
|
|
|
Returns:
|
|
Dict with key "values" mapping to a list of `[R]` value tensors
|
|
per sample.
|
|
"""
|
|
value_list = []
|
|
for logits_chunk, _ in get_responses(
|
|
logits,
|
|
args=args,
|
|
unconcat_tokens=unconcat_tokens,
|
|
total_lengths=total_lengths,
|
|
response_lengths=response_lengths,
|
|
):
|
|
assert logits_chunk.size(-1) == 1, f"{logits_chunk.shape}"
|
|
value_list.append(logits_chunk.squeeze(-1))
|
|
|
|
return {
|
|
"values": value_list,
|
|
}
|
|
|
|
|
|
def compute_advantages_and_returns(args: Namespace, rollout_data: RolloutBatch) -> None:
|
|
"""Compute advantages and returns in-place based on `args.advantage_estimator`.
|
|
|
|
This function extracts rewards, log-probs, values, and masks from
|
|
`rollout_data`, computes KL divergences, then applies the chosen advantage
|
|
estimator. Supported methods: "grpo", "gspo", "ppo", "reinforce_plus_plus",
|
|
and "reinforce_plus_plus_baseline". When `args.normalize_advantages` is
|
|
True, advantages are whitened across the data-parallel group using masked
|
|
statistics.
|
|
|
|
Early returns if both `log_probs` and `values` are None (intermediate
|
|
pipeline stages).
|
|
|
|
Args:
|
|
args: Configuration specifying estimator type, KL coefficient,
|
|
normalization settings, and other hyperparameters.
|
|
rollout_data: Dict containing input lists ("log_probs", "ref_log_probs",
|
|
"rewards", "values", "response_lengths", "loss_masks",
|
|
"total_lengths"). Modified in-place to add "advantages" and
|
|
"returns" keys, each mapping to lists of tensors per sample.
|
|
"""
|
|
log_probs: list[torch.Tensor] = rollout_data.get("rollout_log_probs" if args.use_rollout_logprobs else "log_probs")
|
|
ref_log_probs: list[torch.Tensor] = rollout_data.get("ref_log_probs")
|
|
rewards: list[float] = rollout_data.get("rewards")
|
|
values: None | list[torch.Tensor] = rollout_data.get("values")
|
|
response_lengths: list[int] = rollout_data.get("response_lengths")
|
|
loss_masks: list[torch.Tensor] = rollout_data.get("loss_masks")
|
|
total_lengths: list[int] = rollout_data.get("total_lengths")
|
|
|
|
# return when not the last pp stage.
|
|
if log_probs is None and values is None:
|
|
return
|
|
|
|
if args.kl_coef == 0 or not log_probs:
|
|
# when kl_coef is 0, we won't compute ref_log_prob
|
|
xs = log_probs if log_probs is not None else values
|
|
kl = [torch.zeros_like(x, dtype=torch.float32, device=x.device) for x in xs]
|
|
else:
|
|
kl = [
|
|
compute_approx_kl(
|
|
log_probs[i],
|
|
ref_log_probs[i],
|
|
kl_loss_type=args.kl_loss_type,
|
|
)
|
|
for i in range(len(log_probs))
|
|
]
|
|
|
|
if args.advantage_estimator in ["grpo", "gspo"]:
|
|
rewards = torch.tensor(rewards, dtype=torch.float32, device=kl[0].device)
|
|
returns = get_grpo_returns(rewards, kl)
|
|
# TODO: is the copy necessary?
|
|
advantages = [r for r in returns]
|
|
|
|
elif args.advantage_estimator == "ppo":
|
|
old_rewards = rewards
|
|
rewards = []
|
|
kl_coef = -args.kl_coef
|
|
cp_rank = mpu.get_context_parallel_rank()
|
|
for reward, k in zip(old_rewards, kl, strict=False):
|
|
k *= kl_coef
|
|
if cp_rank == 0:
|
|
k[-1] += reward
|
|
rewards.append(k)
|
|
advantages, returns = get_advantages_and_returns_batch(
|
|
total_lengths, response_lengths, values, rewards, args.gamma, args.lambd
|
|
)
|
|
|
|
elif args.advantage_estimator == "reinforce_plus_plus":
|
|
rewards = torch.tensor(rewards, dtype=torch.float32, device=kl[0].device)
|
|
returns = get_reinforce_plus_plus_returns(
|
|
rewards=rewards,
|
|
kl=kl,
|
|
loss_masks=loss_masks,
|
|
response_lengths=response_lengths,
|
|
total_lengths=total_lengths,
|
|
kl_coef=args.kl_coef,
|
|
gamma=args.gamma,
|
|
)
|
|
advantages = [r for r in returns]
|
|
|
|
elif args.advantage_estimator == "reinforce_plus_plus_baseline":
|
|
rewards = torch.tensor(rewards, dtype=torch.float32, device=kl[0].device)
|
|
advantages = get_reinforce_plus_plus_baseline_advantages(
|
|
rewards=rewards,
|
|
kl=kl,
|
|
loss_masks=loss_masks,
|
|
kl_coef=args.kl_coef,
|
|
)
|
|
returns = advantages
|
|
|
|
elif args.advantage_estimator == "on_policy_distillation":
|
|
student_log_probs = log_probs
|
|
teacher_log_probs = rollout_data.get("teacher_log_probs")
|
|
response_lengths = rollout_data.get("response_lengths")
|
|
|
|
device = student_log_probs[0].device
|
|
teacher_log_probs = [t_log_prob.to(device=device) for t_log_prob in teacher_log_probs]
|
|
teacher_log_probs = [
|
|
t_log_prob[-response_length:]
|
|
for t_log_prob, response_length in zip(teacher_log_probs, response_lengths, strict=False)
|
|
]
|
|
|
|
advantages = [
|
|
teacher_log_prob - student_log_prob
|
|
for teacher_log_prob, student_log_prob in zip(teacher_log_probs, student_log_probs, strict=False)
|
|
]
|
|
|
|
returns = advantages
|
|
|
|
else:
|
|
raise NotImplementedError(f"advantage_estimator {args.advantage_estimator} is not supported. ")
|
|
|
|
# TODO: OpenRLHF always does advantages normalization but veRL doesn't seem to do it.
|
|
if args.normalize_advantages:
|
|
all_advs = torch.cat(advantages)
|
|
cp_size = mpu.get_context_parallel_world_size()
|
|
if cp_size == 1:
|
|
all_masks = torch.cat(loss_masks)
|
|
else:
|
|
mask_chunks = []
|
|
for i in range(len(advantages)):
|
|
total_len = total_lengths[i]
|
|
response_len = response_lengths[i]
|
|
prompt_len = total_len - response_len
|
|
|
|
_, _, _, token_offsets = get_logits_and_tokens_offset_with_cp(total_len, response_len)
|
|
|
|
# Convert global offsets to response-space offsets
|
|
s0, e0 = token_offsets[0]
|
|
s1, e1 = token_offsets[1]
|
|
res_s0, res_e0 = max(0, s0 - prompt_len), max(0, e0 - prompt_len)
|
|
res_s1, res_e1 = max(0, s1 - prompt_len), max(0, e1 - prompt_len)
|
|
|
|
local_mask_parts = []
|
|
full_mask = loss_masks[i]
|
|
if res_e0 > res_s0:
|
|
local_mask_parts.append(full_mask[res_s0:res_e0])
|
|
if res_e1 > res_s1:
|
|
local_mask_parts.append(full_mask[res_s1:res_e1])
|
|
|
|
# Concatenate the parts to form the final mask chunk for this rank and this sequence
|
|
local_mask_chunk = (
|
|
torch.cat(local_mask_parts)
|
|
if local_mask_parts
|
|
else torch.tensor([], device=all_advs.device, dtype=full_mask.dtype)
|
|
)
|
|
mask_chunks.append(local_mask_chunk)
|
|
|
|
all_masks = torch.cat(mask_chunks)
|
|
|
|
if all_masks.numel() > 0:
|
|
assert (
|
|
all_advs.size() == all_masks.size()
|
|
), f"Shape mismatch before whitening: advantages {all_advs.size()}, masks {all_masks.size()}"
|
|
dp_group = mpu.get_data_parallel_group()
|
|
|
|
whitened_advs_flat = distributed_masked_whiten(
|
|
all_advs,
|
|
all_masks,
|
|
process_group=dp_group,
|
|
shift_mean=True,
|
|
)
|
|
chunk_lengths = [chunk.size(0) for chunk in advantages]
|
|
advantages = list(torch.split(whitened_advs_flat, chunk_lengths))
|
|
|
|
rollout_data["advantages"] = advantages
|
|
rollout_data["returns"] = returns
|
|
|
|
|
|
def vanilla_tis_function(
|
|
args,
|
|
*,
|
|
pg_loss: torch.Tensor,
|
|
train_log_probs: list[torch.Tensor],
|
|
rollout_log_probs: list[torch.Tensor],
|
|
loss_masks: list[torch.Tensor],
|
|
**kwargs: Any,
|
|
) -> tuple[torch.Tensor, list[torch.Tensor], dict[str, torch.Tensor]]:
|
|
rollout_log_probs = torch.cat(rollout_log_probs, dim=0)
|
|
old_log_probs = torch.cat(train_log_probs, dim=0)
|
|
tis = torch.exp(old_log_probs - rollout_log_probs)
|
|
tis_abs = (torch.exp(old_log_probs - rollout_log_probs) - 1).abs()
|
|
tis_weights = torch.clamp(tis, min=args.tis_clip_low, max=args.tis_clip)
|
|
tis_clipfrac = (tis_weights != tis).float()
|
|
metrics = {
|
|
"tis": tis.clone().detach(),
|
|
"tis_clipfrac": tis_clipfrac.clone().detach(),
|
|
"tis_abs": tis_abs.clone().detach(),
|
|
}
|
|
pg_loss = pg_loss * tis_weights
|
|
return pg_loss, loss_masks, metrics
|
|
|
|
|
|
def icepop_function(
|
|
args,
|
|
*,
|
|
pg_loss: torch.Tensor,
|
|
train_log_probs: list[torch.Tensor],
|
|
rollout_log_probs: list[torch.Tensor],
|
|
loss_masks: list[torch.Tensor],
|
|
**kwargs: Any,
|
|
) -> tuple[torch.Tensor, list[torch.Tensor], dict[str, torch.Tensor]]:
|
|
rollout_log_probs = torch.cat(rollout_log_probs, dim=0)
|
|
old_log_probs = torch.cat(train_log_probs, dim=0)
|
|
ice_ratio = torch.exp(old_log_probs - rollout_log_probs)
|
|
ice_abs = (torch.exp(old_log_probs - rollout_log_probs) - 1).abs()
|
|
ice_weight = torch.where(
|
|
(ice_ratio >= args.tis_clip_low) & (ice_ratio <= args.tis_clip), ice_ratio, torch.zeros_like(ice_ratio)
|
|
)
|
|
ice_clipfrac = (ice_weight != ice_ratio).float()
|
|
metrics = {
|
|
"tis": ice_ratio.clone().detach(),
|
|
"tis_clipfrac": ice_clipfrac.clone().detach(),
|
|
"tis_abs": ice_abs.clone().detach(),
|
|
}
|
|
pg_loss = pg_loss * ice_weight
|
|
return pg_loss, loss_masks, metrics
|
|
|
|
|
|
def policy_loss_function(
|
|
args: Namespace,
|
|
batch: RolloutBatch,
|
|
logits: torch.Tensor,
|
|
sum_of_sample_mean: Callable[[torch.Tensor], torch.Tensor],
|
|
) -> tuple[torch.Tensor, dict[str, torch.Tensor]]:
|
|
"""Compute policy loss (PPO/GSPO) and metrics.
|
|
|
|
Computes current log-probabilities and entropy from model logits, then
|
|
calculates PPO-style clipped policy gradient loss. For GSPO, gathers
|
|
full sequences via context-parallel all-gather before computing per-sample
|
|
KL. Optionally applies TIS (Truncated Importance Sampling) correction and
|
|
adds KL loss term if configured.
|
|
|
|
Args:
|
|
args: Configuration controlling advantage estimator, clipping thresholds,
|
|
entropy/KL coefficients, and TIS settings.
|
|
batch: Mini-batch containing "advantages", "log_probs" (old policy),
|
|
"unconcat_tokens", "response_lengths", "total_lengths", "loss_masks",
|
|
and optionally "ref_log_probs" and "rollout_log_probs".
|
|
logits: Policy logits with shape `[1, T, V]`.
|
|
sum_of_sample_mean: Reduction function that averages per-sample values.
|
|
|
|
Returns:
|
|
Tuple of `(loss, metrics)` where `loss` is a scalar tensor and `metrics`
|
|
is a dict containing detached scalars: "loss", "pg_loss",
|
|
"entropy_loss", "pg_clipfrac", "ppo_kl". Additional keys "kl_loss",
|
|
"tis", "ois", "tis_clipfrac" are included when the respective features
|
|
are enabled.
|
|
"""
|
|
advantages = torch.cat(batch["advantages"], dim=0)
|
|
old_log_probs = batch["rollout_log_probs"] if args.use_rollout_logprobs else batch["log_probs"]
|
|
|
|
response_lengths = batch["response_lengths"]
|
|
total_lengths = batch["total_lengths"]
|
|
|
|
log_probs_and_entropy = get_log_probs_and_entropy(
|
|
logits,
|
|
args=args,
|
|
unconcat_tokens=batch["unconcat_tokens"],
|
|
total_lengths=total_lengths,
|
|
response_lengths=response_lengths,
|
|
with_entropy=True,
|
|
)
|
|
|
|
log_probs = log_probs_and_entropy["log_probs"]
|
|
|
|
# Pre-gather log probs if needed by OPSM or GSPO to avoid duplicate gathering
|
|
need_full_log_probs = args.use_opsm or args.advantage_estimator == "gspo"
|
|
|
|
full_log_probs = None
|
|
full_old_log_probs = None
|
|
if need_full_log_probs:
|
|
full_log_probs = [
|
|
all_gather_with_cp(log_prob, total_length, response_length)
|
|
for log_prob, total_length, response_length in zip(
|
|
log_probs, total_lengths, response_lengths, strict=False
|
|
)
|
|
]
|
|
full_old_log_probs = [
|
|
all_gather_with_cp(old_log_prob, total_length, response_length)
|
|
for old_log_prob, total_length, response_length in zip(
|
|
old_log_probs, total_lengths, response_lengths, strict=False
|
|
)
|
|
]
|
|
|
|
# Compute OPSM mask if enabled
|
|
if args.use_opsm:
|
|
opsm_mask, opsm_clipfrac = compute_opsm_mask(
|
|
args=args,
|
|
full_log_probs=full_log_probs,
|
|
full_old_log_probs=full_old_log_probs,
|
|
advantages=batch["advantages"],
|
|
loss_masks=batch["loss_masks"],
|
|
)
|
|
|
|
# Compute KL divergence (GSPO uses sequence-level KL, others use per-token KL)
|
|
if args.advantage_estimator == "gspo":
|
|
ppo_kl = compute_gspo_kl(
|
|
full_log_probs=full_log_probs,
|
|
full_old_log_probs=full_old_log_probs,
|
|
local_log_probs=log_probs,
|
|
loss_masks=batch["loss_masks"],
|
|
)
|
|
old_log_probs = torch.cat(old_log_probs, dim=0)
|
|
log_probs = torch.cat(log_probs, dim=0)
|
|
else:
|
|
old_log_probs = torch.cat(old_log_probs, dim=0)
|
|
log_probs = torch.cat(log_probs, dim=0)
|
|
ppo_kl = old_log_probs - log_probs
|
|
|
|
pg_loss, pg_clipfrac = compute_policy_loss(ppo_kl, advantages, args.eps_clip, args.eps_clip_high)
|
|
|
|
if args.use_opsm:
|
|
pg_loss = pg_loss * opsm_mask
|
|
|
|
# Apply off-policy correction using importance sampling if enabled
|
|
if args.get_mismatch_metrics or args.use_tis:
|
|
# NOTE:
|
|
# `tis_func` may apply rejection-sampling style masking (RS) and return `modified_response_masks`.
|
|
# We rebuild `sum_of_sample_mean` with those masks to correct denominators for loss/backprop.
|
|
#
|
|
# However, mismatch/TIS/RS metrics (e.g., "truncate_fraction") are often defined over the
|
|
# *pre-RS* valid tokens. If we aggregate metrics with `modified_response_masks`, the rejected
|
|
# tokens are excluded from the denominator and the metric can be artificially driven to 0.
|
|
# Keep a copy of the original reducer (based on `batch["loss_masks"]`) for metric aggregation.
|
|
sum_of_sample_mean_for_mismatch_metrics = sum_of_sample_mean
|
|
|
|
assert "rollout_log_probs" in batch, "rollout_log_probs must be provided for TIS"
|
|
|
|
ois = (-ppo_kl).exp()
|
|
tis_kwargs = {
|
|
"args": args,
|
|
"pg_loss": pg_loss,
|
|
"train_log_probs": batch["log_probs"],
|
|
"rollout_log_probs": batch["rollout_log_probs"],
|
|
"loss_masks": batch["loss_masks"],
|
|
"total_lengths": total_lengths,
|
|
"response_lengths": response_lengths,
|
|
}
|
|
|
|
if args.custom_tis_function_path is not None:
|
|
tis_func = load_function(args.custom_tis_function_path)
|
|
else:
|
|
tis_func = vanilla_tis_function
|
|
pg_loss, modified_response_masks, tis_metrics = tis_func(**tis_kwargs)
|
|
|
|
# [decouple IS and rejection] Rebuild sum_of_sample_mean with modified_response_masks for denominator correction
|
|
# modified_response_masks will be sliced with cp in get_sum_of_sample_mean
|
|
sum_of_sample_mean = get_sum_of_sample_mean(
|
|
total_lengths, response_lengths, modified_response_masks, args.calculate_per_token_loss
|
|
)
|
|
|
|
pg_loss = sum_of_sample_mean(pg_loss)
|
|
pg_clipfrac = sum_of_sample_mean(pg_clipfrac)
|
|
ppo_kl = sum_of_sample_mean(ppo_kl)
|
|
|
|
# entropy loss
|
|
entropy = log_probs_and_entropy["entropy"]
|
|
entropy = torch.cat(entropy, dim=0)
|
|
entropy_loss = sum_of_sample_mean(entropy)
|
|
|
|
loss = pg_loss - args.entropy_coef * entropy_loss
|
|
|
|
if args.use_kl_loss:
|
|
ref_log_probs = batch["ref_log_probs"]
|
|
ref_log_probs = torch.cat(ref_log_probs, dim=0)
|
|
importance_ratio = None
|
|
if args.use_unbiased_kl:
|
|
importance_ratio = torch.exp(log_probs - old_log_probs)
|
|
kl = compute_approx_kl(
|
|
log_probs,
|
|
ref_log_probs,
|
|
kl_loss_type=args.kl_loss_type,
|
|
importance_ratio=importance_ratio,
|
|
)
|
|
kl_loss = sum_of_sample_mean(kl)
|
|
|
|
loss = loss + args.kl_loss_coef * kl_loss
|
|
|
|
# make sure the gradient could backprop correctly.
|
|
if log_probs.numel() == 0:
|
|
loss += 0 * logits.sum()
|
|
|
|
train_rollout_logprob_abs_diff = None
|
|
importance_weight_mean = None
|
|
importance_weight_std = None
|
|
if "rollout_log_probs" in batch and batch["rollout_log_probs"]:
|
|
rollout_log_probs = torch.cat(batch["rollout_log_probs"], dim=0)
|
|
train_rollout_logprob_abs_diff = sum_of_sample_mean((old_log_probs - rollout_log_probs).abs())
|
|
iw = torch.exp(log_probs.detach() - rollout_log_probs)
|
|
importance_weight_mean = sum_of_sample_mean(iw)
|
|
importance_weight_std = sum_of_sample_mean((iw - 1).pow(2)).sqrt()
|
|
|
|
reported_loss = {
|
|
"loss": loss.clone().detach(),
|
|
"pg_loss": pg_loss.clone().detach(),
|
|
"entropy_loss": entropy_loss.clone().detach(),
|
|
"pg_clipfrac": pg_clipfrac.clone().detach(),
|
|
"ppo_kl": ppo_kl.clone().detach(),
|
|
}
|
|
|
|
if train_rollout_logprob_abs_diff is not None:
|
|
reported_loss["train_rollout_logprob_abs_diff"] = train_rollout_logprob_abs_diff.clone().detach()
|
|
if importance_weight_mean is not None:
|
|
reported_loss["importance_weight_mean"] = importance_weight_mean.clone().detach()
|
|
reported_loss["importance_weight_std"] = importance_weight_std.clone().detach()
|
|
|
|
if args.use_kl_loss:
|
|
reported_loss["kl_loss"] = kl_loss.clone().detach()
|
|
|
|
if args.get_mismatch_metrics or args.use_tis:
|
|
# Aggregate mismatch/TIS/RS related metrics with the *pre-RS* masks.
|
|
# See comment above where `sum_of_sample_mean_for_mismatch_metrics` is defined.
|
|
reported_loss["ois"] = sum_of_sample_mean_for_mismatch_metrics(ois).clone().detach()
|
|
# Assume all metrics are already cloned and detached
|
|
for metric_key, metric_value in tis_metrics.items():
|
|
key_name = f"{metric_key}"
|
|
reported_loss[key_name] = sum_of_sample_mean_for_mismatch_metrics(metric_value)
|
|
|
|
if args.use_opsm:
|
|
reported_loss["opsm_clipfrac"] = opsm_clipfrac
|
|
|
|
return loss, reported_loss
|
|
|
|
|
|
def value_loss_function(
|
|
args: Namespace,
|
|
batch: RolloutBatch,
|
|
logits: torch.Tensor,
|
|
sum_of_sample_mean: Callable[[torch.Tensor], torch.Tensor],
|
|
) -> tuple[torch.Tensor, dict[str, torch.Tensor]]:
|
|
"""Compute clipped value loss and metrics.
|
|
|
|
Extracts current value predictions from `logits`, compares them against
|
|
stored old values with clipping, and computes the maximum of clipped and
|
|
unclipped squared errors (PPO-style value clipping).
|
|
|
|
Args:
|
|
args: Configuration containing `value_clip` threshold.
|
|
batch: Mini-batch with "values" (old predictions), "returns",
|
|
"unconcat_tokens", "total_lengths", and "response_lengths".
|
|
logits: Value head output with shape `[1, T, 1]`.
|
|
sum_of_sample_mean: Reduction function that averages per-sample values.
|
|
|
|
Returns:
|
|
Tuple of `(loss, metrics)` where `loss` is a scalar tensor and
|
|
`metrics` contains detached scalars "value_loss" and "value_clipfrac".
|
|
"""
|
|
old_values = torch.cat(batch["values"], dim=0)
|
|
|
|
values = get_values(
|
|
logits,
|
|
args=args,
|
|
unconcat_tokens=batch["unconcat_tokens"],
|
|
total_lengths=batch["total_lengths"],
|
|
response_lengths=batch["response_lengths"],
|
|
)
|
|
values = torch.cat([value.flatten() for value in values["values"]], dim=0)
|
|
|
|
returns = torch.cat(batch["returns"], dim=0)
|
|
|
|
values_clipfrac = torch.abs(values - old_values) > args.value_clip
|
|
values_clipped = old_values + (values - old_values).clamp(-args.value_clip, args.value_clip)
|
|
surr1 = (values_clipped - returns) ** 2
|
|
surr2 = (values - returns) ** 2
|
|
loss = torch.max(surr1, surr2)
|
|
|
|
loss = sum_of_sample_mean(loss)
|
|
values_clipfrac = sum_of_sample_mean(values_clipfrac.float())
|
|
|
|
# make sure the gradient could backprop correctly.
|
|
if values.numel() == 0:
|
|
loss += 0 * values.sum()
|
|
|
|
reported_loss = {
|
|
"value_loss": loss.clone().detach(),
|
|
"value_clipfrac": values_clipfrac.clone().detach(),
|
|
}
|
|
|
|
return loss, reported_loss
|
|
|
|
|
|
def loss_function(
|
|
args: Namespace,
|
|
batch: RolloutBatch,
|
|
num_microbatches: int,
|
|
logits: torch.Tensor,
|
|
) -> tuple[torch.Tensor, int | torch.Tensor, dict[str, list[str] | torch.Tensor]]:
|
|
"""Dispatch to the configured loss and rescale for Megatron integration.
|
|
|
|
Selects one of "policy_loss", "value_loss", or a custom loss
|
|
function based on `args.loss_type`, computes the loss and metrics, then
|
|
rescales the loss by micro-batch and parallelism factors to integrate with
|
|
Megatron's gradient accumulation.
|
|
|
|
Args:
|
|
args: Configuration specifying `loss_type`, `calculate_per_token_loss`,
|
|
`global_batch_size`, and optionally `custom_loss_function_path`.
|
|
batch: Mini-batch with "loss_masks", "response_lengths", and other
|
|
keys required by the selected loss function.
|
|
num_microbatches: Number of gradient accumulation steps.
|
|
logits: Model outputs (policy or value head).
|
|
|
|
Returns:
|
|
Tuple of `(scaled_loss, normalizer, logging_dict)` where:
|
|
- `scaled_loss` is the loss tensor (scalar) rescaled for Megatron.
|
|
- `normalizer` is `num_tokens` (scalar tensor) if
|
|
`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, ...]).
|
|
"""
|
|
num_tokens = sum([torch.clamp_min(loss_mask.sum(), 1) for loss_mask in batch["loss_masks"]])
|
|
num_samples = len(batch["response_lengths"])
|
|
|
|
sum_of_sample_mean = get_sum_of_sample_mean(
|
|
batch["total_lengths"],
|
|
batch["response_lengths"],
|
|
batch["loss_masks"],
|
|
args.calculate_per_token_loss,
|
|
)
|
|
|
|
loss_type = args.loss_type
|
|
|
|
match loss_type:
|
|
case "policy_loss":
|
|
func = policy_loss_function
|
|
case "value_loss":
|
|
func = value_loss_function
|
|
case "custom_loss":
|
|
func = load_function(args.custom_loss_function_path)
|
|
case _:
|
|
raise ValueError(f"Unknown loss type: {loss_type}")
|
|
|
|
if args.recompute_loss_function:
|
|
loss, log = checkpoint(func, args, batch, logits, sum_of_sample_mean)
|
|
else:
|
|
loss, log = func(args, batch, logits, sum_of_sample_mean)
|
|
|
|
# Here we need to divide by cp_size because to cancel the multiply in Megatron.
|
|
if not args.calculate_per_token_loss:
|
|
loss = (
|
|
loss
|
|
* num_microbatches
|
|
/ args.global_batch_size
|
|
* mpu.get_data_parallel_world_size(with_context_parallel=True)
|
|
)
|
|
else:
|
|
loss = loss * mpu.get_context_parallel_world_size()
|
|
|
|
return (
|
|
loss,
|
|
torch.tensor(num_tokens if args.calculate_per_token_loss else 1, device=logits.device),
|
|
{
|
|
"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,
|
|
),
|
|
},
|
|
)
|