初始化项目，由ModelHub XC社区提供模型
Model: ayh015/myLightningOPD Source: Original Platform
This commit is contained in:
ModelHub XC
230
slime/rollout/data_source.py
Normal file
230
slime/rollout/data_source.py
Normal file
@@ -0,0 +1,230 @@
|
||||
# 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
|
||||
Reference in New Issue
Block a user