600 lines
25 KiB
Python
600 lines
25 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 Sequence
|
|
|
|
import numpy as np
|
|
import torch
|
|
import torch.distributed as dist
|
|
import torch.nn.functional as F
|
|
from megatron.core import mpu
|
|
from megatron.core.packed_seq_params import PackedSeqParams
|
|
|
|
from slime.utils import train_metric_utils
|
|
from slime.utils.data import get_minimum_num_micro_batch_size
|
|
from slime.utils.flops_utils import calculate_fwd_flops
|
|
from slime.utils.metric_utils import compute_pass_rate, compute_rollout_step
|
|
from slime.utils.seqlen_balancing import get_seqlen_balanced_partitions
|
|
from slime.utils.types import RolloutBatch
|
|
|
|
from ...utils import tracking_utils
|
|
from .cp_utils import get_sum_of_sample_mean, slice_with_cp
|
|
|
|
logger = logging.getLogger(__name__)
|
|
|
|
|
|
def get_batch(
|
|
data_iterator: "DataIterator",
|
|
keys: Sequence[str],
|
|
pad_multiplier: int = 128,
|
|
) -> dict[str, torch.Tensor | PackedSeqParams | list[torch.Tensor] | None]:
|
|
"""
|
|
Generate a CP-ready micro-batch with packed sequence parameters.
|
|
|
|
Steps:
|
|
- Fetch raw fields via iterator.
|
|
- Save original token tensors under "unconcat_tokens".
|
|
- Slice tokens into two chunks for Context Parallelism (CP), concatenate, and pad to a configurable multiple.
|
|
- Build cu_seqlens and `PackedSeqParams` with T-H-D layout (T: sequence length, H: attention heads, D: head dimension).
|
|
|
|
Args:
|
|
data_iterator: Iterator providing micro-batch data.
|
|
keys: List of keys to fetch from the iterator.
|
|
pad_multiplier: Multiplier for padding size calculation (default: 128).
|
|
|
|
Returns a dict including:
|
|
- "tokens": torch.LongTensor of shape [1, T_padded] on the current CUDA device
|
|
- "unconcat_tokens": list[torch.LongTensor] for the micro-batch before CP slicing/concat
|
|
- "packed_seq_params": PackedSeqParams with T-H-D settings (cu_seqlens on CUDA, dtype=int)
|
|
Plus any other requested keys forwarded from the iterator.
|
|
"""
|
|
|
|
assert "tokens" in keys
|
|
batch = data_iterator.get_next(keys)
|
|
|
|
tokens = batch["tokens"]
|
|
# use 0 as the pad token id should be fine?
|
|
pad_token_id = 0
|
|
|
|
# for cp, we need all tokens to calculate logprob
|
|
batch["unconcat_tokens"] = tokens
|
|
|
|
cp_size = mpu.get_context_parallel_world_size()
|
|
tokens = [slice_with_cp(t, pad_token_id) for t in tokens]
|
|
|
|
cu_seqlens = [0]
|
|
for t in tokens:
|
|
cu_seqlens.append(cu_seqlens[-1] + t.size(0))
|
|
|
|
tokens = torch.cat(tokens)
|
|
|
|
# Always pad to reduce memory fragmentation and maybe make the computation faster
|
|
pad_size = mpu.get_tensor_model_parallel_world_size() * pad_multiplier
|
|
pad = (pad_size - tokens.size(0) % pad_size) % pad_size
|
|
if pad != 0:
|
|
tokens = F.pad(tokens, (0, pad), value=pad_token_id)
|
|
cu_seqlens.append(cu_seqlens[-1] + pad)
|
|
|
|
# thd requires the cu_seqlens to be of the origin length
|
|
cu_seqlens = torch.tensor(cu_seqlens, dtype=torch.int).cuda() * cp_size
|
|
max_seqlen = (cu_seqlens[1:] - cu_seqlens[:-1]).max().item()
|
|
|
|
packed_seq_params = PackedSeqParams(
|
|
cu_seqlens_q=cu_seqlens,
|
|
cu_seqlens_kv=cu_seqlens,
|
|
max_seqlen_q=max_seqlen,
|
|
max_seqlen_kv=max_seqlen,
|
|
qkv_format="thd",
|
|
)
|
|
|
|
tokens = tokens.unsqueeze(0)
|
|
batch["tokens"] = tokens
|
|
batch["packed_seq_params"] = packed_seq_params
|
|
|
|
# loss masks
|
|
loss_masks = []
|
|
for loss_mask, total_length, response_length in zip(
|
|
batch["loss_masks"],
|
|
batch["total_lengths"],
|
|
batch["response_lengths"],
|
|
strict=True,
|
|
):
|
|
prompt_length = total_length - response_length
|
|
loss_mask = F.pad(loss_mask, (prompt_length - 1, 1), value=0)
|
|
loss_mask = slice_with_cp(loss_mask, 0)
|
|
loss_masks.append(loss_mask)
|
|
loss_masks = torch.cat(loss_masks)
|
|
loss_masks = F.pad(loss_masks, (0, pad), value=0).unsqueeze(0)
|
|
assert loss_masks.shape == tokens.shape, f"loss_masks.shape: {loss_masks.shape}, tokens.shape: {tokens.shape}"
|
|
batch["full_loss_masks"] = loss_masks
|
|
|
|
# Process multimodal training tensors if present
|
|
multimodal_train_inputs = batch.get("multimodal_train_inputs", None)
|
|
if multimodal_train_inputs is not None:
|
|
multimodal_data = {} # key -> concatenated tensor
|
|
multimodal_num_items = {} # key -> list of item counts per sequence
|
|
for mm_input_dict in multimodal_train_inputs:
|
|
if mm_input_dict is not None:
|
|
for key, mm_tensor in mm_input_dict.items():
|
|
if key not in multimodal_data:
|
|
multimodal_data[key] = mm_tensor
|
|
multimodal_num_items[key] = [mm_tensor.size(0)]
|
|
else:
|
|
multimodal_data[key] = torch.cat([multimodal_data[key], mm_tensor], dim=0)
|
|
multimodal_num_items[key].append(mm_tensor.size(0))
|
|
batch["multimodal_train_inputs"] = multimodal_data
|
|
batch["multimodal_num_items"] = multimodal_num_items
|
|
|
|
return batch
|
|
|
|
|
|
def gather_log_data(
|
|
metric_name: str,
|
|
args: Namespace,
|
|
rollout_id: int,
|
|
log_dict: dict[str, float],
|
|
) -> dict[str, float] | None:
|
|
"""
|
|
Gather per-rank metrics, reduce by mean on the DP source rank, and log.
|
|
|
|
Expects `log_dict` to contain plain scalars. The DP source rank prints and
|
|
optionally logs to WandB/TensorBoard with a step derived from `rollout_id` and
|
|
batch sizes. Returns the reduced dict on the DP source rank; returns None on others.
|
|
"""
|
|
|
|
if mpu.get_data_parallel_rank(with_context_parallel=True) == 0:
|
|
dp_size = mpu.get_data_parallel_world_size(with_context_parallel=True)
|
|
|
|
gathered_log_dict = [None] * dp_size
|
|
# Not sure if this will be a performance bottleneck.
|
|
dist.gather_object(
|
|
log_dict,
|
|
gathered_log_dict,
|
|
dst=mpu.get_data_parallel_src_rank(with_context_parallel=True),
|
|
group=mpu.get_data_parallel_group_gloo(with_context_parallel=True),
|
|
)
|
|
|
|
reduced_log_dict = {
|
|
f"{metric_name}/{key}": sum([d[key] for d in gathered_log_dict]) / dp_size for key in log_dict
|
|
}
|
|
logger.info(f"{metric_name} {rollout_id}: {reduced_log_dict}")
|
|
|
|
# Calculate step once to avoid duplication
|
|
step = compute_rollout_step(args, rollout_id)
|
|
reduced_log_dict["rollout/step"] = step
|
|
tracking_utils.log(args, reduced_log_dict, step_key="rollout/step")
|
|
|
|
return reduced_log_dict
|
|
else:
|
|
dist.gather_object(
|
|
log_dict,
|
|
None,
|
|
dst=mpu.get_data_parallel_src_rank(with_context_parallel=True),
|
|
group=mpu.get_data_parallel_group_gloo(with_context_parallel=True),
|
|
)
|
|
return None
|
|
|
|
|
|
class DataIterator:
|
|
"""Micro-batch iterator over rollout dicts.
|
|
|
|
Supports either fixed contiguous micro-batches or an explicit per-step
|
|
index schedule (for dynamic batch sizing / sequence-length balancing).
|
|
"""
|
|
|
|
def __init__(
|
|
self,
|
|
rollout_data: RolloutBatch,
|
|
micro_batch_size: int | None = None,
|
|
micro_batch_indices: list[list[int]] | None = None,
|
|
) -> None:
|
|
"""Initialize an iterator over `rollout_data`.
|
|
|
|
Args:
|
|
rollout_data: Dict of per-sample fields for the local step.
|
|
micro_batch_size: Fixed contiguous slice size when not using dynamic scheduling.
|
|
micro_batch_indices: Explicit indices per micro-batch when using dynamic balancing.
|
|
Must be mutually exclusive with `micro_batch_size`.
|
|
"""
|
|
self.rollout_data = rollout_data
|
|
self.micro_batch_size = micro_batch_size
|
|
self.micro_batch_indices = micro_batch_indices
|
|
assert micro_batch_size is None or micro_batch_indices is None
|
|
self.offset = 0
|
|
|
|
# Keys that are batch-level (not per-sample) and should be passed through as-is
|
|
BATCH_LEVEL_KEYS = set()
|
|
|
|
def get_next(self, keys: Sequence[str]) -> dict[str, list[object] | None]:
|
|
"""Return the next micro-batch for the requested keys.
|
|
|
|
- If `micro_batch_indices` is provided, selects rows according to the current
|
|
index list for each requested key.
|
|
- Otherwise, slices a contiguous window of size `micro_batch_size` starting
|
|
at the current offset.
|
|
|
|
Returns a dict mapping each key to a list subset (or None if absent).
|
|
"""
|
|
batch = {}
|
|
for key in keys:
|
|
vals = self.rollout_data.get(key, None)
|
|
if vals is None:
|
|
batch[key] = None
|
|
elif key in self.BATCH_LEVEL_KEYS:
|
|
# Batch-level keys are not per-sample, pass through as-is
|
|
batch[key] = vals
|
|
else:
|
|
if self.micro_batch_indices is not None:
|
|
indices = self.micro_batch_indices[self.offset]
|
|
batch[key] = [vals[i] for i in indices]
|
|
else:
|
|
assert self.offset + self.micro_batch_size <= len(
|
|
vals
|
|
), f"offset: {self.offset}, micro_batch_size: {self.micro_batch_size}, len(vals): {len(vals)}"
|
|
batch[key] = vals[self.offset : self.offset + self.micro_batch_size]
|
|
|
|
if self.micro_batch_indices is not None:
|
|
self.offset += 1
|
|
else:
|
|
self.offset += self.micro_batch_size
|
|
return batch
|
|
|
|
def reset(self) -> "DataIterator":
|
|
"""Reset internal offset to the start and return self."""
|
|
self.offset = 0
|
|
return self
|
|
|
|
|
|
def get_data_iterator(
|
|
args: Namespace,
|
|
model: torch.nn.Module | Sequence[torch.nn.Module],
|
|
rollout_data: RolloutBatch,
|
|
) -> tuple[list[DataIterator], list[int]]:
|
|
"""
|
|
Create iterators and a micro-batch schedule for a rollout step.
|
|
|
|
- If `use_dynamic_batch_size` is False, splits into fixed-size contiguous
|
|
micro-batches of `micro_batch_size`.
|
|
- If True, computes the number of micro-batches per local step based on
|
|
`max_tokens_per_gpu` and per-sample lengths, all-reduces to a DP-wide
|
|
maximum, optionally enforces divisibility for Virtual Pipeline Parallelism (VPP), and builds a balanced
|
|
index schedule to equalize token counts across micro-batches.
|
|
|
|
Returns `(data_iterators, num_microbatches)` where:
|
|
- `data_iterators`: list of `DataIterator`, one per VPP stage (size 1 if VPP disabled)
|
|
- `num_microbatches`: list[int], one per local step in the rollout (length = steps)
|
|
"""
|
|
dp_size = mpu.get_data_parallel_world_size(with_context_parallel=False)
|
|
dp_group = mpu.get_data_parallel_group()
|
|
vpp_size = mpu.get_virtual_pipeline_model_parallel_world_size()
|
|
if vpp_size is None:
|
|
vpp_size = 1
|
|
if vpp_size > 1:
|
|
from megatron.core.utils import get_model_config
|
|
|
|
config = get_model_config(model[0])
|
|
microbatch_group_size_per_vp_stage = config.microbatch_group_size_per_vp_stage
|
|
cp_size = mpu.get_context_parallel_world_size()
|
|
|
|
num_local_samples = len(rollout_data["total_lengths"])
|
|
num_local_gbs = args.global_batch_size // dp_size
|
|
num_steps_per_rollout = num_local_samples // num_local_gbs
|
|
|
|
def _generate_data_iterator(rollout_data, micro_batch_size, micro_batch_indices=None):
|
|
data_iterator = []
|
|
for _ in range(vpp_size):
|
|
data_iterator.append(DataIterator(rollout_data, micro_batch_size, micro_batch_indices))
|
|
return data_iterator
|
|
|
|
if not args.use_dynamic_batch_size:
|
|
num_microbatches = [num_local_gbs // args.micro_batch_size for _ in range(num_steps_per_rollout)]
|
|
data_iterator = _generate_data_iterator(rollout_data, args.micro_batch_size)
|
|
else:
|
|
assert args.max_tokens_per_gpu is not None
|
|
# calculate the number of mirobatches for each step
|
|
samples = rollout_data["total_lengths"]
|
|
assert len(samples) == num_local_samples
|
|
num_microbatches = []
|
|
for i in range(num_steps_per_rollout):
|
|
start, end = i * num_local_gbs, (i + 1) * num_local_gbs
|
|
num_microbatches.append(
|
|
get_minimum_num_micro_batch_size(samples[start:end], args.max_tokens_per_gpu * cp_size)
|
|
)
|
|
|
|
num_microbatches = torch.tensor(num_microbatches, dtype=torch.int, device=torch.cuda.current_device())
|
|
dist.all_reduce(num_microbatches, op=dist.ReduceOp.MAX, group=dp_group)
|
|
|
|
if vpp_size > 1:
|
|
# vpp requies the number of microbatches to be divisible by vpp_size
|
|
num_microbatches = torch.clamp(
|
|
num_microbatches // microbatch_group_size_per_vp_stage * microbatch_group_size_per_vp_stage,
|
|
min=1,
|
|
)
|
|
|
|
num_microbatches = num_microbatches.tolist()
|
|
|
|
# balance the each micro batch
|
|
samples = rollout_data["total_lengths"]
|
|
# balance the number of mirobatches across steps
|
|
micro_batch_indices = []
|
|
for i, num_mbs in enumerate(num_microbatches):
|
|
start, end = i * num_local_gbs, (i + 1) * num_local_gbs
|
|
samples = rollout_data["total_lengths"][start:end]
|
|
partitions = get_seqlen_balanced_partitions(samples, num_mbs, equal_size=False)
|
|
for j in range(num_mbs):
|
|
for k in range(len(partitions[j])):
|
|
partitions[j][k] += start
|
|
micro_batch_indices.extend(partitions)
|
|
|
|
assert len(set(sum(micro_batch_indices, []))) == num_local_samples
|
|
|
|
data_iterator = _generate_data_iterator(rollout_data, None, micro_batch_indices)
|
|
|
|
return (
|
|
data_iterator,
|
|
num_microbatches,
|
|
)
|
|
|
|
|
|
def log_rollout_data(rollout_id: int, args: Namespace, rollout_data: RolloutBatch) -> None:
|
|
"""
|
|
Summarize rollout fields and log reduced metrics on PP last stage, TP rank 0.
|
|
|
|
- Tensor-valued lists are concatenated and averaged. For token-level metrics
|
|
like log-probs/returns/advantages/values, computes a CP-correct sample mean
|
|
using `loss_masks` and total/response lengths.
|
|
- Non-tensor lists are averaged elementwise.
|
|
- Scalars are converted to Python numbers.
|
|
"""
|
|
if mpu.get_tensor_model_parallel_rank() == 0 and mpu.is_pipeline_last_stage():
|
|
cp_size = mpu.get_context_parallel_world_size()
|
|
log_dict = {}
|
|
response_lengths = rollout_data["response_lengths"]
|
|
loss_masks = rollout_data["loss_masks"]
|
|
total_lengths = rollout_data["total_lengths"]
|
|
|
|
for key, val in rollout_data.items():
|
|
if key in [
|
|
"tokens",
|
|
"multimodal_train_inputs",
|
|
"loss_masks",
|
|
"sample_indices",
|
|
"rollout_routed_experts",
|
|
]:
|
|
continue
|
|
# Skip None values
|
|
if val is None:
|
|
continue
|
|
# Upload per sample mean for each rollout value
|
|
# There are the following assumptions:
|
|
# - Each dp rank has the same number of samples
|
|
if isinstance(val, (list, tuple)):
|
|
# Filter out None entries before processing.
|
|
val = [v for v in val if v is not None]
|
|
if not val:
|
|
continue
|
|
if all(isinstance(v, torch.Tensor) for v in val):
|
|
# NOTE: Here we have to do the clone().detach(), otherwise the tensor will be
|
|
# modified in place and will cause problem for the next rollout.
|
|
val = torch.cat(val).clone().detach()
|
|
if key in ["log_probs", "ref_log_probs", "rollout_log_probs", "returns", "advantages", "values"]:
|
|
sum_of_sample_mean = get_sum_of_sample_mean(total_lengths, response_lengths, loss_masks)
|
|
val = cp_size * sum_of_sample_mean(val) / len(loss_masks)
|
|
else:
|
|
val = val.mean() * cp_size
|
|
else:
|
|
# Mixed Tensor/scalar list.
|
|
# Convert everything to float scalar for logging.
|
|
val = sum(float(v.mean()) if isinstance(v, torch.Tensor) else float(v) for v in val) / len(val)
|
|
elif isinstance(val, torch.Tensor):
|
|
val = val.float().mean()
|
|
else:
|
|
raise ValueError(f"Unsupported type: {type(val)} for key: {key}")
|
|
log_dict[key] = val.item() if isinstance(val, torch.Tensor) else val
|
|
|
|
reduced_log_dict = gather_log_data("rollout", args, rollout_id, log_dict)
|
|
if args.ci_test and reduced_log_dict is not None:
|
|
if (
|
|
rollout_id == 0
|
|
and "rollout/log_probs" in reduced_log_dict
|
|
and "rollout/ref_log_probs" in reduced_log_dict
|
|
):
|
|
assert reduced_log_dict["rollout/log_probs"] == reduced_log_dict["rollout/ref_log_probs"]
|
|
if "rollout/log_probs" in reduced_log_dict:
|
|
assert -0.5 < reduced_log_dict["rollout/log_probs"] < 0
|
|
if "rollout/entropy" in reduced_log_dict:
|
|
assert 0 < reduced_log_dict["rollout/entropy"] < 0.5
|
|
|
|
if args.log_multi_turn:
|
|
log_multi_turn_data(rollout_id, args, rollout_data)
|
|
if args.log_passrate:
|
|
log_passrate(rollout_id, args, rollout_data)
|
|
|
|
if args.log_correct_samples:
|
|
if mpu.get_tensor_model_parallel_rank() == 0 and mpu.is_pipeline_last_stage():
|
|
cp_size = mpu.get_context_parallel_world_size()
|
|
log_dict = {}
|
|
response_lengths = rollout_data["response_lengths"]
|
|
loss_masks = rollout_data["loss_masks"]
|
|
total_lengths = rollout_data["total_lengths"]
|
|
|
|
def quantile(total_value, n_quantiles, data) -> dict:
|
|
import math
|
|
|
|
assert n_quantiles > 1, f"n_quantiles({n_quantiles}) must be greater than 1."
|
|
|
|
quantiles = [((i + 1) / n_quantiles) for i in range(n_quantiles)]
|
|
cut_points = [total_value * q for q in quantiles]
|
|
cut_points[-1] = total_value
|
|
|
|
count = [0] * n_quantiles
|
|
for d in data:
|
|
for i, point in enumerate(cut_points):
|
|
if d <= point:
|
|
count[i] += 1
|
|
break
|
|
|
|
total = sum(count) + 1e-9
|
|
percentile = [c / total for c in count]
|
|
|
|
percentile = {f"p{min(math.ceil(q*100),100)}": p for q, p in zip(quantiles, percentile, strict=True)}
|
|
return percentile
|
|
|
|
raw_rewards = rollout_data["raw_reward"]
|
|
# Additional metrics for correct cases are calculated separately below.
|
|
correct_response_lengths = []
|
|
correct_total_lengths = []
|
|
correct_loss_masks = []
|
|
correct_entropy = []
|
|
for i, raw_reward in enumerate(raw_rewards):
|
|
if raw_reward == 1:
|
|
correct_response_lengths.append(response_lengths[i])
|
|
correct_total_lengths.append(total_lengths[i])
|
|
correct_loss_masks.append(loss_masks[i])
|
|
correct_entropy.append(-rollout_data["log_probs"][i])
|
|
num_correct_responses = len(correct_total_lengths)
|
|
rollout_data["correct_response_lengths"] = correct_response_lengths
|
|
correct_response_length_percentile = quantile(
|
|
args.rollout_max_response_len, 4, rollout_data["correct_response_lengths"]
|
|
)
|
|
for p, val in correct_response_length_percentile.items():
|
|
rollout_data[f"correct_length/{p}"] = [val] * num_correct_responses
|
|
if len(correct_entropy) > 0:
|
|
sum_of_sample_mean = get_sum_of_sample_mean(
|
|
correct_total_lengths, correct_response_lengths, correct_loss_masks
|
|
)
|
|
correct_entropy = sum_of_sample_mean(torch.cat(correct_entropy, dim=0))
|
|
rollout_data["correct_entropy"] = [correct_entropy.item()] * num_correct_responses
|
|
else:
|
|
rollout_data["correct_entropy"] = [0] * num_correct_responses
|
|
|
|
|
|
def log_multi_turn_data(rollout_id: int, args: Namespace, rollout_data: RolloutBatch) -> None:
|
|
"""
|
|
Log multi-turn auxiliary metrics such as raw/observed response lengths and rounds.
|
|
|
|
Operates only on PP last stage and TP rank 0. Uses GPU tensors when available
|
|
to compute statistics without host transfers.
|
|
"""
|
|
if mpu.get_tensor_model_parallel_rank() == 0 and mpu.is_pipeline_last_stage():
|
|
log_dict = {}
|
|
for key, val in rollout_data.items():
|
|
if key == "loss_masks":
|
|
if val: # Check if val is not empty
|
|
device = val[0].device # Get device from first tensor
|
|
|
|
# Vectorized length calculation using torch
|
|
raw_response_lengths = torch.tensor([v.shape[0] for v in val], dtype=torch.float32, device=device)
|
|
log_dict["raw_response_length/response_length_mean"] = raw_response_lengths.mean().item()
|
|
log_dict["raw_response_length/response_length_max"] = raw_response_lengths.max().item()
|
|
log_dict["raw_response_length/response_length_min"] = raw_response_lengths.min().item()
|
|
log_dict["raw_response_length/response_length_clip_ratio"] = (
|
|
(raw_response_lengths >= args.rollout_max_response_len).float().mean().item()
|
|
)
|
|
|
|
# Vectorized sum calculation using torch - stay on GPU
|
|
wo_obs_response_lengths = torch.tensor(
|
|
[v.sum().item() for v in val], dtype=torch.float32, device=device
|
|
)
|
|
log_dict["wo_obs_response_length/response_length_mean"] = wo_obs_response_lengths.mean().item()
|
|
log_dict["wo_obs_response_length/response_length_max"] = wo_obs_response_lengths.max().item()
|
|
log_dict["wo_obs_response_length/response_length_min"] = wo_obs_response_lengths.min().item()
|
|
if key == "round_number":
|
|
# Use numpy for vectorized round number statistics
|
|
round_number_array = np.array(val)
|
|
log_dict["multi_turn_metric/round_number_mean"] = np.mean(round_number_array)
|
|
log_dict["multi_turn_metric/round_number_max"] = np.max(round_number_array)
|
|
log_dict["multi_turn_metric/round_number_min"] = np.min(round_number_array)
|
|
gather_log_data("multi_turn", args, rollout_id, log_dict)
|
|
|
|
|
|
def log_passrate(rollout_id: int, args: Namespace, rollout_data: RolloutBatch) -> None:
|
|
"""
|
|
Compute pass@k metrics from `raw_reward` groups and log the results.
|
|
|
|
`raw_reward` is reshaped to `[group_number, group_size]`, then pass@k is
|
|
estimated per problem and averaged.
|
|
"""
|
|
if mpu.get_tensor_model_parallel_rank() == 0 and mpu.is_pipeline_last_stage():
|
|
log_dict = {}
|
|
for key, val in rollout_data.items():
|
|
if key != "raw_reward":
|
|
continue
|
|
|
|
log_dict |= compute_pass_rate(
|
|
flat_rewards=val,
|
|
group_size=args.n_samples_per_prompt,
|
|
num_groups=args.rollout_batch_size,
|
|
)
|
|
|
|
gather_log_data("passrate", args, rollout_id, log_dict)
|
|
|
|
|
|
def log_perf_data(rollout_id: int, args: Namespace) -> None:
|
|
train_metric_utils.log_perf_data_raw(
|
|
rollout_id=rollout_id,
|
|
args=args,
|
|
is_primary_rank=(
|
|
mpu.get_tensor_model_parallel_rank() == 0
|
|
and mpu.is_pipeline_last_stage()
|
|
and mpu.get_data_parallel_rank(with_context_parallel=True) == 0
|
|
),
|
|
compute_total_fwd_flops=lambda seq_lens: calculate_fwd_flops(seqlens=seq_lens, args=args)
|
|
/ dist.get_world_size()
|
|
/ 1e12,
|
|
)
|
|
|
|
|
|
def sync_actor_critic_data(
|
|
args: Namespace,
|
|
rollout_data: RolloutBatch | None = None,
|
|
group: dist.ProcessGroup | None = None,
|
|
) -> None:
|
|
"""
|
|
Broadcast `values` (from critic) and optionally `log_probs`/`ref_log_probs`
|
|
(from actor) across PP ranks to align data dependencies.
|
|
|
|
- Values are broadcast from src=1.
|
|
- Log-probs and ref-log-probs are broadcast from src=0 when KL is used.
|
|
Updates `rollout_data` in place with the synchronized tensors.
|
|
"""
|
|
log_probs_key = "log_probs" if not args.use_rollout_logprobs else "rollout_log_probs"
|
|
values, log_probs, ref_log_probs = map(rollout_data.get, ("values", log_probs_key, "ref_log_probs"))
|
|
|
|
# return when not the pp last stage
|
|
if not values and not log_probs:
|
|
return
|
|
|
|
handles = []
|
|
|
|
if not values:
|
|
values = [torch.empty_like(log_prob) for log_prob in log_probs]
|
|
for value in values:
|
|
handles.append(dist.broadcast(value, src=1, group=group, async_op=True))
|
|
|
|
if args.kl_coef != 0 or args.use_kl_loss:
|
|
if not log_probs:
|
|
log_probs = [torch.empty_like(value) for value in values]
|
|
if not ref_log_probs:
|
|
ref_log_probs = [torch.empty_like(value) for value in values]
|
|
for ref_log_prob, log_prob in zip(ref_log_probs, log_probs, strict=False):
|
|
handles.append(dist.broadcast(log_prob, src=0, group=group, async_op=True))
|
|
handles.append(dist.broadcast(ref_log_prob, src=0, group=group, async_op=True))
|
|
|
|
for handle in handles:
|
|
handle.wait()
|
|
|
|
rollout_data.update(
|
|
{
|
|
k: v
|
|
for k, v in {
|
|
"values": values,
|
|
log_probs_key: log_probs,
|
|
"ref_log_probs": ref_log_probs,
|
|
}.items()
|
|
if v is not None
|
|
}
|
|
)
|