""" Data processing utilities for GRPO trainer. Handles data retrieval from Atropos API, padding, batching, and advantage normalization. Also extracts inference logprobs for proper GRPO loss computation: - Inference logprobs serve as π_old (reference policy) for importance sampling - They are batched and padded to align token-by-token with training labels """ import json import math import time from typing import List, Optional, Tuple import numpy as np import torch from .api import get_batch def pad_data_to_good_offset( data: dict, batch_size: int, extract_inference_logprobs: bool = True, ) -> Tuple[ List[torch.Tensor], # token_batches List[torch.Tensor], # label_batches List[torch.Tensor], # advantage_batches List[torch.Tensor], # temperature_batches Optional[List[torch.Tensor]], # inference_logprob_batches (aligned with labels) ]: """ Pad and batch data from the Atropos API. Processes raw batch data into properly padded tensors suitable for training: - Pads token sequences to nearest multiple of 64 - Normalizes advantage scores - Extracts temperature values - Extracts and pads inference logprobs for proper GRPO loss computation Args: data: Raw batch data from Atropos API batch_size: Size of each training batch extract_inference_logprobs: Whether to extract inference logprobs Returns: Tuple of (token_batches, label_batches, advantage_batches, temperature_batches, inference_logprob_batches) inference_logprob_batches is None if extract_inference_logprobs=False or no logprobs in data Note: inference_logprob_batches are padded with 0.0 at positions where labels == -100. This allows token-by-token alignment during GRPO loss computation. """ max_token_len = max( [max([len(x) for x in item["tokens"]]) for item in data["batch"]] ) # Pad to nearest multiple of 64 for GPU efficiency good_multiple = 64 if (max_token_len - 1) % (good_multiple) != 0: max_token_len = math.ceil((max_token_len - 1) / (good_multiple)) * good_multiple token_setup_len = max_token_len + 1 # +1 for causal shift else: token_setup_len = max_token_len max_token_len = max_token_len - 1 # -1 for causal shift # Process all items input_ids = [] labels = [] advantages = [] lengths = [] temperatures = [] inference_logprobs_padded: List[np.ndarray] = [] # Padded to match labels shape has_any_logprobs = False for item in data["batch"]: # Normalize advantage scores scores = np.array(item["scores"]) if len(scores) > 1: scores = scores - scores.mean() scores = scores / max(scores.std(), 1e-8) item["scores"] = scores # Handle score overrides if item["overrides"] is not None: for i in range(len(item["overrides"])): if item["overrides"][i].get("set_advantage_to_zero", False): item["scores"][i] = 0 # Process each sample in the item for i in range(len(item["tokens"])): seq_len = len(item["tokens"][i]) lengths.append( math.ceil((seq_len - 1) / good_multiple) * good_multiple ) # Create labels with padding (-100 for masked positions) label_item = np.concatenate([ np.array(item["masks"][i]), np.full( max(0, token_setup_len - seq_len), -100, dtype=np.int32, ), ]) # Pad tokens item["tokens"][i] = np.concatenate([ np.array(item["tokens"][i]), np.zeros( max(0, token_setup_len - seq_len), dtype=np.int32, ), ]) input_ids.append(item["tokens"][i][:-1]) # Remove last for causal labels.append(label_item[1:]) # Shift by 1 for causal advantages.append(item["scores"][i]) # Extract and pad inference logprobs to match labels shape # Inference logprobs are ONLY for generated tokens (where labels != -100) # We need to create a padded array that aligns position-by-position if extract_inference_logprobs and "inference_logprobs" in item: if i < len(item["inference_logprobs"]): raw_logprobs = np.array(item["inference_logprobs"][i], dtype=np.float32) has_any_logprobs = True # Create padded logprobs array matching label_item shape # Fill with 0.0 (will be masked out during loss computation) padded_logprobs = np.zeros(token_setup_len, dtype=np.float32) # The inference logprobs correspond to generated tokens # Find positions where labels != -100 (generated positions) mask_arr = np.array(item["masks"][i]) generated_positions = np.where(mask_arr != -100)[0] # Fill in inference logprobs at generated positions n_to_fill = min(len(raw_logprobs), len(generated_positions)) if n_to_fill > 0: padded_logprobs[generated_positions[:n_to_fill]] = raw_logprobs[:n_to_fill] # Shift by 1 to match causal label shift inference_logprobs_padded.append(padded_logprobs[1:]) else: # No logprobs for this sample, use zeros inference_logprobs_padded.append(np.zeros(token_setup_len - 1, dtype=np.float32)) elif extract_inference_logprobs: # No inference_logprobs in item, use zeros inference_logprobs_padded.append(np.zeros(token_setup_len - 1, dtype=np.float32)) # Extract temperature (priority: override > generation_params > group_overrides > 1.0) t = 1.0 if ( item.get("overrides") and i < len(item["overrides"]) and isinstance(item["overrides"][i], dict) and ("temperature" in item["overrides"][i]) ): t = float(item["overrides"][i]["temperature"]) elif item.get("generation_params") and ("temperature" in item["generation_params"]): t = float(item["generation_params"]["temperature"]) elif item.get("group_overrides") and ("temperature" in item["group_overrides"]): t = float(item["group_overrides"]["temperature"]) temperatures.append(t) # Batch the data token_batches = [] label_batches = [] advantage_batches = [] temperature_batches = [] inference_logprob_batches = [] for i in range(len(input_ids) // batch_size): start = i * batch_size end = (i + 1) * batch_size token_batches.append( torch.tensor(np.stack(input_ids[start:end], axis=0)) ) label_batches.append( torch.tensor(np.stack(labels[start:end], axis=0)) ) advantage_batches.append( torch.tensor(np.stack(advantages[start:end], axis=0)).view(-1, 1) ) temperature_batches.append( torch.tensor( np.array(temperatures[start:end], dtype=np.float32) ).view(-1, 1, 1) ) # Batch inference logprobs (same shape as labels) if extract_inference_logprobs and inference_logprobs_padded: inference_logprob_batches.append( torch.tensor(np.stack(inference_logprobs_padded[start:end], axis=0)) ) # Return inference logprob batches if we have any real logprobs final_logprob_batches = inference_logprob_batches if (has_any_logprobs and inference_logprob_batches) else None return token_batches, label_batches, advantage_batches, temperature_batches, final_logprob_batches def get_data( batch_size: int, seq_len: int, atropos_url: str = "http://localhost:8000", extract_inference_logprobs: bool = True, ) -> Tuple[ List[Tuple[ List[torch.Tensor], # token_batches List[torch.Tensor], # label_batches List[torch.Tensor], # advantage_batches List[torch.Tensor], # temperature_batches Optional[List[torch.Tensor]], # inference_logprob_batches ]], None, # Legacy return (no longer used) ]: """ Fetch and process training data from the Atropos API. Continuously polls the API until data is available, then processes all available batches. Args: batch_size: Size of each training batch seq_len: Maximum sequence length (for reference, not used directly) atropos_url: URL of the Atropos API server extract_inference_logprobs: Whether to extract inference logprobs for GRPO loss Returns: Tuple of (batches, None) - batches: List of processed batch tuples, each containing: (token_batches, label_batches, advantage_batches, temperature_batches, inference_logprob_batches) - inference_logprob_batches are aligned with labels for proper GRPO loss computation """ batches = [] while True: data = get_batch(url=atropos_url) if data["batch"] is not None: # Save batch for debugging with open("temp.json", "w", encoding="utf-8") as f: json.dump(data, f) # Process and accumulate batches (now includes batched inference logprobs) token_batches, label_batches, adv_batches, temp_batches, inf_logprob_batches = \ pad_data_to_good_offset(data, batch_size, extract_inference_logprobs) # Include inference logprob batches in the tuple batches.append((token_batches, label_batches, adv_batches, temp_batches, inf_logprob_batches)) elif len(batches) > 0: # Return accumulated batches when no more data return batches, None else: # Wait for data time.sleep(1)