slime/backends/fsdp_utils/data_packing.py

# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
# SPDX-License-Identifier: Apache-2.0

"""Data packing utilities for FSDP backend to reduce padding overhead."""

import math

import torch
import torch.nn.functional as F

from slime.utils.seqlen_balancing import get_seqlen_balanced_partitions


def pack_sequences(
    tokens: list[list[int]],
    loss_masks: list[list[int]],
    rewards: list[float],
    raw_rewards: list,
    response_lengths: list[int],
    advantages: list[float],
    returns: list[float],
    rollout_log_probs: list[list[float]] | None = None,
    multimodal_train_inputs: list[dict] | None = None,
    max_tokens_per_gpu: int | None = None,
    num_packs: int | None = None,
) -> list[dict]:
    """
    Pack sequences into dense batches with cumulative sequence lengths.

    Args:
        tokens: List of token sequences
        loss_masks: List of loss masks
        rewards: List of rewards per sequence
        raw_rewards: List of raw rewards per sequence
        response_lengths: List of response lengths per sequence
        advantages: List of advantages per sequence
        returns: List of returns per sequence
        rollout_log_probs: List of rollout log probabilities per sequence
        multimodal_train_inputs: List of dict of multimodal tensors for training per sequence
        max_tokens_per_gpu: Maximum tokens per GPU pack
        num_packs: Explicit number of packs to create

    Returns:
        List of packed batches with tokens, masks, cu_seqlens, rewards, raw_rewards, response_lengths, advantages, returns
    """
    if not tokens:
        return []

    seq_lengths = [len(t) for t in tokens]

    # Determine number of packs and use balanced partitioning
    if num_packs:
        k_partitions = num_packs
    elif max_tokens_per_gpu:
        total_tokens = sum(seq_lengths)
        k_partitions = max(1, math.ceil(total_tokens / max_tokens_per_gpu))
    else:
        k_partitions = 1

    # Use balanced partitioning for optimal load distribution
    partitions = get_seqlen_balanced_partitions(
        seq_lengths, k_partitions=k_partitions, equal_size=False  # Allow variable sizes for better balance
    )

    # Pack each partition
    result = []
    for indices in partitions:
        # Build cumulative sequence lengths
        cu_seqlens = [0]
        flat_tokens = []
        flat_masks = []
        flat_positionids = []
        flat_advantages = []
        flat_returns = []
        flat_rollout_log_probs = []

        for i in indices:
            seq_tokens = tokens[i]
            seq_mask = loss_masks[i]
            seq_positionids = list(range(len(seq_tokens)))

            flat_tokens.extend(seq_tokens)
            flat_positionids.extend(seq_positionids)
            flat_masks.extend(seq_mask)
            flat_advantages.extend(advantages[i])
            flat_returns.extend(returns[i])
            if rollout_log_probs:
                flat_rollout_log_probs.extend(rollout_log_probs[i])
            cu_seqlens.append(cu_seqlens[-1] + len(seq_tokens))

        packed_batch = {
            "tokens": torch.tensor(flat_tokens, dtype=torch.long),
            "loss_masks": torch.tensor(flat_masks, dtype=torch.int),
            "position_ids": torch.tensor(flat_positionids, dtype=torch.int),
            "cu_seqlens": torch.tensor(cu_seqlens, dtype=torch.int32),
            "rewards": torch.tensor([rewards[i] for i in indices], dtype=torch.float32),
            "raw_reward": [raw_rewards[i] for i in indices],
            "response_lengths": [response_lengths[i] for i in indices],
            "advantages": torch.tensor(flat_advantages, dtype=torch.float32),
            "returns": torch.tensor(flat_returns, dtype=torch.float32),
            "rollout_log_probs": torch.tensor(
                flat_rollout_log_probs, dtype=torch.float32, device=torch.cuda.current_device()
            ),
        }

        # Collect and add multimodal training tensors for this partition
        if multimodal_train_inputs:
            multimodal_data = {}  # key -> concatenated tensor
            multimodal_num_items = {}  # key -> list of item counts per sequence
            for i in indices:
                for key, mm_tensor in multimodal_train_inputs[i].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))
            packed_batch["multimodal_train_inputs"] = multimodal_data
            packed_batch["multimodal_num_items"] = multimodal_num_items

        result.append(packed_batch)

    return result


def unpack_sequences(packed_batch: dict) -> list[dict]:
    """
    Unpack sequences from a packed batch.

    Args:
        packed_batch: Packed batch

    Returns:
        List of unpacked batches
    """

    cu_seqlens = packed_batch["cu_seqlens"]
    num_sequences = len(cu_seqlens) - 1
    response_lengths = packed_batch["response_lengths"]
    multimodal_num_items = packed_batch.get("multimodal_num_items", {})

    instances = []

    # Calculate pad_length by counting trailing zeros
    tokens = packed_batch["tokens"]
    nonzero_indices = (tokens != 0).nonzero(as_tuple=True)[0]
    if len(nonzero_indices) > 0:
        # Last non-zero index, pad_length is everything after it
        pad_length = len(tokens) - nonzero_indices[-1].item() - 1
    else:
        pad_length = 0  # No padding if no non-zero tokens (or all zeros)
    for i in range(num_sequences):
        start_idx = cu_seqlens[i].item()
        end_idx = cu_seqlens[i + 1].item()
        instance = {}

        # Copy any additional attributes that might exist in the packed batch
        for key, value in packed_batch.items():
            if key not in instance:
                # Skip multimodal_num_items - it's metadata
                if key == "multimodal_num_items":
                    continue
                # Handle multimodal_train_inputs dict: split each tensor using multimodal_num_items
                elif key == "multimodal_train_inputs" and isinstance(value, dict):
                    instance[key] = {}
                    for mm_key, mm_tensor in value.items():
                        if mm_key in multimodal_num_items:
                            num_items_list = multimodal_num_items[mm_key]
                            start_mm_idx = sum(num_items_list[:i])
                            end_mm_idx = start_mm_idx + num_items_list[i]
                            if num_items_list[i] > 0:
                                instance[key][mm_key] = mm_tensor[start_mm_idx:end_mm_idx]
                # For tensor attributes, we need to slice them appropriately
                elif isinstance(value, torch.Tensor):
                    if key in ["log_probs", "ref_log_probs", "cur_log_probs", "entropy"]:
                        # These are computed from logits[:-1] so they have length seq_len-1
                        instance[key] = value[
                            end_idx - 1 - response_lengths[i] - pad_length : end_idx - 1 - pad_length
                        ]
                    elif key == "rollout_log_probs":
                        # rollout_log_probs is packed based on response_lengths, so slice differently
                        instance[key] = value[sum(response_lengths[:i]) : sum(response_lengths[: i + 1])]
                    elif key in ["tokens", "position_ids"]:
                        # For other tensor attributes, try to slice them
                        if len(value) > start_idx:
                            instance[key] = value[start_idx:end_idx]
                        else:
                            raise ValueError(f"Attribute {key} is not found in the packed batch")
                    elif key in ["loss_masks", "advantages", "returns"]:
                        instance[key] = value[sum(response_lengths[:i]) : sum(response_lengths[: i + 1])]
                elif isinstance(value, list):
                    instance[key] = value[i]
                else:
                    raise ValueError(f"Attribute {key} is not found in the packed batch")

        instances.append(instance)

    return instances


def pad_packed_sequence_with_cp(packed_sequence: dict, cp_size: int) -> dict:
    """Pad packed sequence to make total length divisible by cp_size.

    Args:
        packed_sequence: Packed sequence dict containing tokens, position_ids, cu_seqlens, etc.
        cp_size: Context parallelism world size

    Returns:
        Padded packed sequence
    """
    seq_length = len(packed_sequence["tokens"])
    # Calculate padding needed: (cp_size - seq_length % cp_size) % cp_size
    remainder = seq_length % cp_size
    pad_length = (cp_size - remainder) % cp_size

    if pad_length > 0:
        packed_sequence["tokens"] = F.pad(packed_sequence["tokens"], (0, pad_length), value=0)
        packed_sequence["position_ids"] = F.pad(packed_sequence["position_ids"], (0, pad_length), value=0)
        packed_sequence["loss_masks"] = F.pad(packed_sequence["loss_masks"], (0, pad_length), value=0)
        packed_sequence["cu_seqlens"][-1] += pad_length
    return packed_sequence