File size: 8,268 Bytes
6011e08 | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 117 118 119 120 121 122 123 124 125 126 127 128 129 130 131 132 133 134 135 136 137 138 139 140 141 142 143 144 145 146 147 148 149 150 151 152 153 154 155 156 157 158 159 160 161 162 163 164 165 166 167 168 169 170 171 172 173 174 175 176 177 178 179 180 181 182 183 184 185 186 187 188 189 190 191 192 193 194 195 196 197 198 199 200 201 202 203 204 205 206 207 208 209 210 211 212 213 214 215 216 217 218 219 220 221 222 223 224 225 226 227 228 229 230 231 | # SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
# SPDX-License-Identifier: Apache-2.0
import abc
import copy
import logging
import os
from pathlib import Path
import torch
from slime.utils.data import create_dataset
from slime.utils.misc import load_function
from slime.utils.processing_utils import load_processor, load_tokenizer
from slime.utils.types import Sample
logger = logging.getLogger(__name__)
class DataSource(abc.ABC):
@abc.abstractmethod
def get_samples(self, num_samples: int) -> list[list[Sample]]:
"""
Return num_samples samples
"""
@abc.abstractmethod
def add_samples(self, samples: list[list[Sample]]):
"""
Add samples to the data source
"""
@abc.abstractmethod
def save(self, rollout_id):
"""
Save the state of the data source
"""
@abc.abstractmethod
def load(self, rollout_id=None):
"""
Load the state of the data source
"""
# TODO may further refactor data-loading part later
class RolloutDataSource(DataSource):
def __init__(self, args):
self.args = args
self.epoch_id = 0
self.sample_group_index = 0
self.sample_index = 0
self.sample_offset = 0
# TODO remove this
self.metadata = {}
if args.rollout_global_dataset:
tokenizer = load_tokenizer(args.hf_checkpoint, trust_remote_code=True)
processor = load_processor(args.hf_checkpoint, trust_remote_code=True)
# TODO move (during the refactor)
if (d := args.dump_details) is not None:
tokenizer.save_pretrained(Path(d) / "tokenizer")
if processor:
processor.save_pretrained(Path(d) / "processor")
self.dataset = create_dataset(
args.prompt_data,
tokenizer=tokenizer,
processor=processor,
max_length=args.rollout_max_prompt_len,
prompt_key=args.input_key,
multimodal_keys=args.multimodal_keys,
label_key=args.label_key,
metadata_key=args.metadata_key,
tool_key=args.tool_key,
apply_chat_template=args.apply_chat_template,
apply_chat_template_kwargs=args.apply_chat_template_kwargs,
seed=args.rollout_seed,
)
if self.args.rollout_shuffle:
self.dataset.shuffle(self.epoch_id)
else:
self.dataset = None
def get_samples(self, num_samples):
# TODO further improve code
if self.dataset is not None:
if self.sample_offset + num_samples <= len(self.dataset):
prompt_samples = self.dataset.samples[self.sample_offset : self.sample_offset + num_samples]
self.sample_offset += num_samples
else:
prompt_samples = self.dataset.samples[self.sample_offset :]
num_samples -= len(prompt_samples)
self.epoch_id += 1
if self.args.rollout_shuffle:
self.dataset.shuffle(self.epoch_id)
prompt_samples += self.dataset.samples[:num_samples]
self.sample_offset = num_samples
else:
prompt_samples = [Sample() for _ in range(num_samples)]
samples = []
for prompt_sample in prompt_samples:
group = []
for _ in range(self.args.n_samples_per_prompt):
sample = copy.deepcopy(prompt_sample)
sample.group_index = self.sample_group_index
sample.index = self.sample_index
self.sample_index += 1
group.append(sample)
self.sample_group_index += 1
samples.append(group)
return samples
def add_samples(self, samples: list[list[Sample]]):
raise RuntimeError(f"Cannot add samples to {self.__class__.__name__}. This is a read-only data source.")
def save(self, rollout_id):
if not self.args.rollout_global_dataset:
return
state_dict = {
"sample_offset": self.sample_offset,
"epoch_id": self.epoch_id,
"sample_group_index": self.sample_group_index,
"sample_index": self.sample_index,
"metadata": self.metadata,
# Save wandb_run_id for resume support
"wandb_run_id": getattr(self.args, "wandb_run_id", None),
}
path = os.path.join(self.args.save, f"rollout/global_dataset_state_dict_{rollout_id}.pt")
os.makedirs(os.path.dirname(path), exist_ok=True)
torch.save(state_dict, path)
def load(self, rollout_id=None):
if not self.args.rollout_global_dataset:
return
if self.args.load is None:
return
path = os.path.join(self.args.load, f"rollout/global_dataset_state_dict_{rollout_id}.pt")
if not os.path.exists(path):
logger.info(f"Checkpoint {path} does not exist.")
return
logger.info(f"load metadata from {path}")
logger.info(f"load metadata: {self.metadata}")
state_dict = torch.load(path)
self.sample_offset = state_dict.get("sample_offset", 0)
self.epoch_id = state_dict.get("epoch_id", 0)
self.sample_group_index = state_dict.get("sample_group_index", 0)
self.sample_index = state_dict.get("sample_index", 0)
self.metadata = state_dict.get("metadata", {})
# Load wandb_run_id for resume support (only if not already set)
if not getattr(self.args, "wandb_run_id", None):
loaded_wandb_run_id = state_dict.get("wandb_run_id")
if loaded_wandb_run_id:
self.args.wandb_run_id = loaded_wandb_run_id
logger.info(f"Loaded wandb_run_id from checkpoint: {loaded_wandb_run_id}")
if self.args.rollout_global_dataset and self.args.rollout_shuffle:
self.dataset.shuffle(self.epoch_id)
class RolloutDataSourceWithBuffer(RolloutDataSource):
def __init__(self, args):
super().__init__(args)
self.buffer = []
if self.args.buffer_filter_path is None:
self.buffer_filter = pop_first
else:
self.buffer_filter = load_function(self.args.buffer_filter_path)
def get_samples(self, num_samples: int) -> list[list[Sample]]:
"""
Return num_samples samples
"""
samples = self._get_samples_from_buffer(num_samples)
num_samples -= len(samples)
if num_samples == 0:
return samples
samples += super().get_samples(num_samples=num_samples)
return samples
def _get_samples_from_buffer(self, num_samples: int) -> list[list[Sample]]:
if len(self.buffer) == 0 or num_samples == 0:
return []
samples = self.buffer_filter(self.args, None, self.buffer, num_samples)
return samples
def add_samples(self, samples: list[list[Sample]]):
"""
Add a sample group to buffer.
"""
if not samples:
return
assert isinstance(samples, list), f"samples must be a list, got {type(samples)}"
assert isinstance(samples[0], list), f"the elements of samples must be list, got {type(samples[0])}"
for i in range(0, len(samples)):
assert (
len(samples[i]) == self.args.n_samples_per_prompt
), f"the length of the elements of samples must be equal to n_samples_per_prompt, got {len(samples[i])} != {self.args.n_samples_per_prompt}"
group = samples[i] # type: ignore
self.buffer.append(group)
# TODO remove
def update_metadata(self, metadata: dict):
self.metadata.update(metadata)
# TODO remove
def get_metadata(self):
return self.metadata
def get_buffer_length(self):
return len(self.buffer)
def pop_first(args, rollout_id, buffer: list[list[Sample]], num_samples: int) -> list[list[Sample]]:
num_to_pop = min(len(buffer), num_samples)
samples = buffer[:num_to_pop]
del buffer[:num_to_pop]
return samples
|