231 lines
8.1 KiB
Python
231 lines
8.1 KiB
Python
# 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
|