chore: import upstream snapshot with attribution
This commit is contained in:
wehub-resource-sync
@@ -0,0 +1,5 @@
|
||||
# Copyright (c) Microsoft Corporation.
|
||||
# SPDX-License-Identifier: Apache-2.0
|
||||
|
||||
# DeepSpeed Team
|
||||
'''Copyright The Microsoft DeepSpeed Team'''
|
||||
@@ -0,0 +1,887 @@
|
||||
# Copyright (c) Microsoft Corporation.
|
||||
# SPDX-License-Identifier: Apache-2.0
|
||||
|
||||
# DeepSpeed Team
|
||||
|
||||
import glob
|
||||
import os
|
||||
import sys
|
||||
from collections import defaultdict
|
||||
import csv
|
||||
import time
|
||||
from multiprocessing import Process, Manager
|
||||
import numpy as np
|
||||
import torch
|
||||
from torch.utils.data import BatchSampler, SequentialSampler, DataLoader, Subset
|
||||
|
||||
import deepspeed.comm as dist
|
||||
from deepspeed.utils import logger
|
||||
from deepspeed.runtime.data_pipeline.data_sampling.indexed_dataset import MMapIndexedDataset, valid_dtypes
|
||||
from deepspeed.runtime.data_pipeline.data_sampling.utils import split_dataset, split_index, create_mmap_dataset_builder, close_mmap_dataset_builder, find_fit_int_dtype
|
||||
|
||||
|
||||
class DataAnalyzer(object):
|
||||
|
||||
def __init__(self,
|
||||
dataset,
|
||||
num_workers=1,
|
||||
worker_id=0,
|
||||
num_threads=1,
|
||||
num_threads_reduce=1,
|
||||
specific_threads=[],
|
||||
batch_size=1,
|
||||
metric_names=[],
|
||||
metric_functions=[],
|
||||
metric_types=[],
|
||||
metric_dtypes=[],
|
||||
save_path="./",
|
||||
collate_fn=None,
|
||||
custom_map_init=None,
|
||||
custom_map_update=None,
|
||||
custom_map_finalize=None,
|
||||
custom_reduce=None,
|
||||
sample_indices=None):
|
||||
super().__init__()
|
||||
self.dataset = dataset
|
||||
self.num_workers = num_workers
|
||||
self.worker_id = worker_id
|
||||
self.num_threads = num_threads
|
||||
self.num_threads_reduce = num_threads_reduce
|
||||
self.specific_threads = specific_threads
|
||||
self.batch_size = batch_size
|
||||
self.metric_names = metric_names
|
||||
self.metric_functions = metric_functions
|
||||
self.metric_types = metric_types
|
||||
self.metric_dtypes = metric_dtypes
|
||||
self.save_path = save_path
|
||||
self.collate_fn = collate_fn
|
||||
self.custom_map_init = custom_map_init
|
||||
self.custom_map_update = custom_map_update
|
||||
self.custom_map_finalize = custom_map_finalize
|
||||
self.custom_reduce = custom_reduce
|
||||
self.sample_indices = sample_indices
|
||||
|
||||
def init_metric_results(self, thread_id, metric_names, metric_types, metric_dtypes, save_path, worker_id):
|
||||
metric_results = []
|
||||
for m_idx in range(len(metric_names)):
|
||||
metric_name, metric_type, metric_dtype = metric_names[m_idx], \
|
||||
metric_types[m_idx], metric_dtypes[m_idx]
|
||||
assert metric_dtype in valid_dtypes, f"metric_dtype {metric_dtype} not supported. Supported dtypes {valid_dtypes}"
|
||||
metric_save_path = f"{save_path}/{metric_name}/worker{worker_id}_thread{thread_id}/"
|
||||
os.makedirs(metric_save_path, exist_ok=True)
|
||||
if metric_type == 'single_value_per_sample':
|
||||
sample_to_metric_fname = f"{metric_save_path}/{metric_name}_sample_to_metric"
|
||||
sample_to_metric_builder = create_mmap_dataset_builder(sample_to_metric_fname, metric_dtype)
|
||||
metric_to_sample_fname = f"{metric_save_path}/{metric_name}_metric_to_sample"
|
||||
for _f in glob.glob(f"{glob.escape(metric_to_sample_fname)}*"):
|
||||
os.remove(_f)
|
||||
metric_to_sample_dict = defaultdict(list)
|
||||
metric_results.append({
|
||||
"sample_to_metric_fname": sample_to_metric_fname,
|
||||
"sample_to_metric_builder": sample_to_metric_builder,
|
||||
"metric_to_sample_fname": metric_to_sample_fname,
|
||||
"metric_to_sample_dict": metric_to_sample_dict
|
||||
})
|
||||
elif metric_type == 'accumulate_value_over_samples':
|
||||
metric_value = None
|
||||
metric_value_fname = f"{metric_save_path}/{metric_name}_metric_value"
|
||||
metric_results.append({"metric_value": metric_value, "metric_value_fname": metric_value_fname})
|
||||
return metric_results
|
||||
|
||||
def update_metric_results(self,
|
||||
data,
|
||||
metric_types,
|
||||
metric_dtypes,
|
||||
metric_functions,
|
||||
metric_results,
|
||||
batch_start_idx=0):
|
||||
for m_idx in range(len(metric_types)):
|
||||
metric_type, metric_dtype, metric_function, metric_result = metric_types[m_idx], \
|
||||
metric_dtypes[m_idx], metric_functions[m_idx], metric_results[m_idx]
|
||||
metric_values = metric_function(data)
|
||||
|
||||
assert torch.is_tensor(metric_values) or isinstance(metric_values, np.ndarray), \
|
||||
"metric_function must return a tensor or array"
|
||||
assert metric_values.dtype == metric_dtype, \
|
||||
f"metric_function result dtype {metric_values.dtype} does not match metric_dtype {metric_dtype}"
|
||||
if isinstance(metric_values, np.ndarray):
|
||||
metric_values = torch.from_numpy(metric_values)
|
||||
|
||||
if metric_type == 'single_value_per_sample':
|
||||
for row in range(metric_values.size()[0]):
|
||||
sample_idx = batch_start_idx + row # sample idx following dataset iteration order
|
||||
if isinstance(data, dict) and 'index' in data: # Megatron use case, idx provided in 'index' field
|
||||
sample_idx = data['index'][row][0].item()
|
||||
elif self.sample_indices is not None: # user defined shuffling of indices
|
||||
sample_idx = self.sample_indices[sample_idx]
|
||||
metric_result["sample_to_metric_builder"].add_item(metric_values[row].reshape(-1))
|
||||
metric_result["metric_to_sample_dict"][metric_values[row].item()].append(sample_idx)
|
||||
for m_value in metric_result["metric_to_sample_dict"]:
|
||||
if len(metric_result["metric_to_sample_dict"][m_value]) > 100:
|
||||
metric_fname = metric_result["metric_to_sample_fname"]
|
||||
with open(f"{metric_fname}_{m_value}.csv", 'a') as f:
|
||||
writer = csv.writer(f)
|
||||
writer.writerows([metric_result["metric_to_sample_dict"][m_value]])
|
||||
metric_result["metric_to_sample_dict"][m_value] = []
|
||||
elif metric_type == 'accumulate_value_over_samples':
|
||||
if metric_result["metric_value"] is None:
|
||||
metric_result["metric_value"] = metric_values
|
||||
else:
|
||||
metric_result["metric_value"].add_(metric_values)
|
||||
|
||||
def finalize_metric_results(self, metric_types, metric_dtypes, metric_results):
|
||||
for m_idx in range(len(metric_types)):
|
||||
metric_type, metric_dtype, metric_result = metric_types[m_idx], \
|
||||
metric_dtypes[m_idx], metric_results[m_idx]
|
||||
if metric_type == 'single_value_per_sample':
|
||||
metric_fname = metric_result["sample_to_metric_fname"]
|
||||
close_mmap_dataset_builder(metric_result["sample_to_metric_builder"], metric_fname)
|
||||
for m_value in metric_result["metric_to_sample_dict"]:
|
||||
if len(metric_result["metric_to_sample_dict"][m_value]) > 0:
|
||||
metric_fname = metric_result["metric_to_sample_fname"]
|
||||
with open(f"{metric_fname}_{m_value}.csv", 'a') as f:
|
||||
writer = csv.writer(f)
|
||||
writer.writerows([metric_result["metric_to_sample_dict"][m_value]])
|
||||
metric_result["metric_to_sample_dict"][m_value] = []
|
||||
elif metric_type == 'accumulate_value_over_samples':
|
||||
if metric_result["metric_value"] is not None:
|
||||
metric_value_builder = create_mmap_dataset_builder(metric_result["metric_value_fname"],
|
||||
metric_dtype)
|
||||
metric_value_builder.add_item(metric_result["metric_value"].reshape(-1))
|
||||
close_mmap_dataset_builder(metric_value_builder, metric_result["metric_value_fname"])
|
||||
|
||||
def run_map_helper(self, thread_id):
|
||||
start_idx, end_idx = self.thread_splits[thread_id][0], \
|
||||
self.thread_splits[thread_id][1]
|
||||
logger.info(f"worker {self.worker_id} thread {thread_id}: start working " \
|
||||
f"on data subset {start_idx} to {end_idx}")
|
||||
thread_dataset = Subset(self.dataset, list(range(start_idx, end_idx)))
|
||||
sampler = BatchSampler(SequentialSampler(thread_dataset), batch_size=self.batch_size, drop_last=False)
|
||||
iterator = iter(
|
||||
DataLoader(thread_dataset,
|
||||
batch_sampler=sampler,
|
||||
num_workers=0,
|
||||
collate_fn=self.collate_fn,
|
||||
pin_memory=False))
|
||||
if self.custom_map_init is None:
|
||||
metric_results = self.init_metric_results(thread_id, self.metric_names, self.metric_types,
|
||||
self.metric_dtypes, self.save_path, self.worker_id)
|
||||
else:
|
||||
metric_results = self.custom_map_init(thread_id, self.metric_names, self.metric_types, self.metric_dtypes,
|
||||
self.save_path, self.worker_id)
|
||||
total_sample = len(thread_dataset)
|
||||
processed_sample = 0
|
||||
start = time.time()
|
||||
while True:
|
||||
try:
|
||||
data = next(iterator)
|
||||
batch_start_idx = start_idx + processed_sample
|
||||
if self.custom_map_update is None:
|
||||
self.update_metric_results(data, self.metric_types, self.metric_dtypes, self.metric_functions,
|
||||
metric_results, batch_start_idx)
|
||||
else:
|
||||
self.custom_map_update(data, self.metric_types, self.metric_dtypes, self.metric_functions,
|
||||
metric_results, batch_start_idx)
|
||||
processed_sample += len(data)
|
||||
duration = (time.time() - start) / 3600.0
|
||||
remain_duration = duration * total_sample / processed_sample - duration
|
||||
logger.info(
|
||||
f"worker {self.worker_id} thread {thread_id}: {processed_sample} " \
|
||||
f"out of {total_sample} processed in {duration:.2f} hr, " \
|
||||
f"estimated to finish in {remain_duration:.2f} hr")
|
||||
except StopIteration:
|
||||
logger.info(f"worker {self.worker_id} thread {thread_id}: reach end of file")
|
||||
break
|
||||
if self.custom_map_finalize is None:
|
||||
self.finalize_metric_results(self.metric_types, self.metric_dtypes, metric_results)
|
||||
else:
|
||||
self.custom_map_finalize(self.metric_types, self.metric_dtypes, metric_results)
|
||||
logger.info(f"worker {self.worker_id} thread {thread_id}: finished")
|
||||
|
||||
def run_map(self):
|
||||
self.worker_splits, self.thread_splits = split_dataset(self.dataset, self.num_workers, self.worker_id,
|
||||
self.num_threads)
|
||||
if len(self.specific_threads) > 0:
|
||||
threads_to_run = self.specific_threads
|
||||
else:
|
||||
threads_to_run = list(range(self.num_threads))
|
||||
if self.num_threads > 1:
|
||||
p = []
|
||||
for thread in threads_to_run:
|
||||
p.append(Process(target=self.run_map_helper, args=(thread, )))
|
||||
p[thread].start()
|
||||
|
||||
for thread in threads_to_run:
|
||||
p[thread].join()
|
||||
else:
|
||||
assert self.num_threads == 1
|
||||
self.run_map_helper(0)
|
||||
|
||||
def get_metric_value_percentiles(self, metric_name, num_sample_per_value, total_num_samples):
|
||||
logger.info(f"Checking the value percentiles of metric {metric_name}...")
|
||||
processed_samples = 0
|
||||
current_percentile = 5
|
||||
for key in sorted(num_sample_per_value.keys()):
|
||||
processed_samples += num_sample_per_value[key]
|
||||
if processed_samples >= total_num_samples * current_percentile / 100.0:
|
||||
logger.info(f"Metric {metric_name} {current_percentile}th percentile: {key}")
|
||||
current_percentile += 5
|
||||
|
||||
def merge_gather_map_stats(self, num_workers, num_threads, num_threads_reduce, t_idx_reduce, metric_save_path,
|
||||
metric_name, return_dict):
|
||||
results = []
|
||||
for w_idx in range(num_workers):
|
||||
for t_idx in range(num_threads):
|
||||
if (w_idx * num_threads + t_idx) % num_threads_reduce == t_idx_reduce:
|
||||
w_metric_save_path = f"{metric_save_path}/worker{w_idx}_thread{t_idx}/"
|
||||
w_sample_to_metric_fname = f"{w_metric_save_path}/{metric_name}_sample_to_metric"
|
||||
w_sample_to_metric = MMapIndexedDataset(w_sample_to_metric_fname, skip_warmup=True)
|
||||
unique_v = list(np.unique(w_sample_to_metric))
|
||||
sample_to_metric_count = len(w_sample_to_metric)
|
||||
logger.info(f"Finished gathering map stats from worker {w_idx} thread {t_idx}.")
|
||||
results.append([unique_v, sample_to_metric_count])
|
||||
return_dict[t_idx_reduce] = results
|
||||
|
||||
def merge_sample_to_metric(self, t_idx_reduce, metric_save_path, metric_name, metric_value_dtype,
|
||||
map_worker_thread):
|
||||
sample_to_metric_fname = f"{metric_save_path}/{metric_name}_sample_to_metric_thread{t_idx_reduce}"
|
||||
sample_to_metric_builder = create_mmap_dataset_builder(sample_to_metric_fname, metric_value_dtype)
|
||||
for w_t in map_worker_thread:
|
||||
w_metric_save_path = f"{metric_save_path}/worker{w_t[0]}_thread{w_t[1]}/"
|
||||
w_sample_to_metric_fname = f"{w_metric_save_path}/{metric_name}_sample_to_metric"
|
||||
w_data = MMapIndexedDataset(w_sample_to_metric_fname, skip_warmup=True)
|
||||
for row in range(len(w_data)):
|
||||
sample_to_metric_builder.add_item(torch.tensor(w_data[row].astype(np.int64), dtype=torch.long))
|
||||
logger.info(f"Finished merge_sample_to_metric from worker {w_t[0]} thread {w_t[1]}.")
|
||||
close_mmap_dataset_builder(sample_to_metric_builder, sample_to_metric_fname)
|
||||
|
||||
def merge_metric_to_sample(self, t_idx_reduce, metric_save_path, metric_name, sample_idx_dtype, metric_value_dtype,
|
||||
unique_metric_values, num_workers, num_threads):
|
||||
index_to_sample_fname = f"{metric_save_path}/{metric_name}_index_to_sample_thread{t_idx_reduce}"
|
||||
index_to_sample_builder = create_mmap_dataset_builder(index_to_sample_fname, sample_idx_dtype)
|
||||
index_to_metric_fname = f"{metric_save_path}/{metric_name}_index_to_metric_thread{t_idx_reduce}"
|
||||
index_to_metric_builder = create_mmap_dataset_builder(index_to_metric_fname, metric_value_dtype)
|
||||
for unique_v in unique_metric_values:
|
||||
samples = []
|
||||
for w_idx in range(num_workers):
|
||||
for t_idx in range(num_threads):
|
||||
w_metric_save_path = f"{metric_save_path}/worker{w_idx}_thread{t_idx}/"
|
||||
w_metric_to_sample_fname = f"{w_metric_save_path}/{metric_name}_metric_to_sample_{unique_v}.csv"
|
||||
if os.path.isfile(w_metric_to_sample_fname):
|
||||
with open(w_metric_to_sample_fname, 'r') as f:
|
||||
datareader = csv.reader(f)
|
||||
for row in datareader:
|
||||
samples += [int(x) for x in row]
|
||||
index_to_sample_builder.add_item(torch.tensor(samples, dtype=torch.long))
|
||||
index_to_metric_builder.add_item(torch.tensor([unique_v], dtype=torch.long))
|
||||
logger.info(f"Finished reducing metric {metric_name} value {unique_v}.")
|
||||
close_mmap_dataset_builder(index_to_sample_builder, index_to_sample_fname)
|
||||
close_mmap_dataset_builder(index_to_metric_builder, index_to_metric_fname)
|
||||
|
||||
def merge_map_results(self, dataset, metric_names, metric_types, save_path, num_workers, num_threads,
|
||||
num_threads_reduce):
|
||||
total_num_samples = len(dataset)
|
||||
sample_idx_dtype = find_fit_int_dtype(0, total_num_samples - 1)
|
||||
logger.info(
|
||||
f"Total number of data samples: {total_num_samples}. Will use {sample_idx_dtype} to store the sample indexes."
|
||||
)
|
||||
for m_idx in range(len(metric_names)):
|
||||
metric_name, metric_type = metric_names[m_idx], metric_types[m_idx]
|
||||
if metric_type == 'single_value_per_sample':
|
||||
metric_save_path = f"{save_path}/{metric_name}/"
|
||||
sample_to_metric_count = 0
|
||||
unique_metric_values = set([])
|
||||
manager = Manager()
|
||||
return_dict = manager.dict()
|
||||
p = []
|
||||
for t_idx_reduce in range(num_threads_reduce):
|
||||
p.append(
|
||||
Process(target=self.merge_gather_map_stats,
|
||||
args=(
|
||||
num_workers,
|
||||
num_threads,
|
||||
num_threads_reduce,
|
||||
t_idx_reduce,
|
||||
metric_save_path,
|
||||
metric_name,
|
||||
return_dict,
|
||||
)))
|
||||
p[t_idx_reduce].start()
|
||||
for t_idx_reduce in range(num_threads_reduce):
|
||||
p[t_idx_reduce].join()
|
||||
for t_idx_reduce in range(num_threads_reduce):
|
||||
results = return_dict[t_idx_reduce]
|
||||
for res in results:
|
||||
unique_metric_values = unique_metric_values.union(set(res[0]))
|
||||
sample_to_metric_count += res[1]
|
||||
value_max = max(unique_metric_values)
|
||||
value_min = min(unique_metric_values)
|
||||
assert sample_to_metric_count == total_num_samples, "The number of samples in map result files are not correct. It's possible that some map worker didn't finish successfully."
|
||||
metric_value_dtype = find_fit_int_dtype(value_min, value_max)
|
||||
logger.info(
|
||||
f"Metric {metric_name} has values between {value_min} and {value_max}. Will use {metric_value_dtype} to store the metric values."
|
||||
)
|
||||
|
||||
# sample_to_metric
|
||||
map_worker_thread = []
|
||||
for w_idx in range(num_workers):
|
||||
for t_idx in range(num_threads):
|
||||
map_worker_thread.append([w_idx, t_idx])
|
||||
thread_splits = split_index(0, len(map_worker_thread), num_threads_reduce)
|
||||
p = []
|
||||
for t_idx_reduce in range(num_threads_reduce):
|
||||
start_idx, end_idx = thread_splits[t_idx_reduce][0], thread_splits[t_idx_reduce][1]
|
||||
p.append(
|
||||
Process(target=self.merge_sample_to_metric,
|
||||
args=(
|
||||
t_idx_reduce,
|
||||
metric_save_path,
|
||||
metric_name,
|
||||
metric_value_dtype,
|
||||
map_worker_thread[start_idx:end_idx],
|
||||
)))
|
||||
p[t_idx_reduce].start()
|
||||
for t_idx_reduce in range(num_threads_reduce):
|
||||
p[t_idx_reduce].join()
|
||||
|
||||
sample_to_metric_fname = f"{metric_save_path}/{metric_name}_sample_to_metric"
|
||||
sample_to_metric_builder = create_mmap_dataset_builder(sample_to_metric_fname, metric_value_dtype)
|
||||
for t_idx_reduce in range(num_threads_reduce):
|
||||
chunk_fname = f"{metric_save_path}/{metric_name}_sample_to_metric_thread{t_idx_reduce}"
|
||||
logger.info(f"Merging file {chunk_fname}")
|
||||
sample_to_metric_builder.merge_file_(chunk_fname)
|
||||
close_mmap_dataset_builder(sample_to_metric_builder, sample_to_metric_fname)
|
||||
sample_to_metric = MMapIndexedDataset(sample_to_metric_fname, skip_warmup=True)
|
||||
assert len(sample_to_metric) == total_num_samples
|
||||
|
||||
# metric_to_sample
|
||||
unique_metric_values = list(sorted(unique_metric_values))
|
||||
thread_splits = split_index(0, len(unique_metric_values), num_threads_reduce)
|
||||
p = []
|
||||
for t_idx_reduce in range(num_threads_reduce):
|
||||
start_idx, end_idx = thread_splits[t_idx_reduce][0], thread_splits[t_idx_reduce][1]
|
||||
p.append(
|
||||
Process(target=self.merge_metric_to_sample,
|
||||
args=(
|
||||
t_idx_reduce,
|
||||
metric_save_path,
|
||||
metric_name,
|
||||
sample_idx_dtype,
|
||||
metric_value_dtype,
|
||||
unique_metric_values[start_idx:end_idx],
|
||||
num_workers,
|
||||
num_threads,
|
||||
)))
|
||||
p[t_idx_reduce].start()
|
||||
for t_idx_reduce in range(num_threads_reduce):
|
||||
p[t_idx_reduce].join()
|
||||
index_to_sample_fname = f"{metric_save_path}/{metric_name}_index_to_sample"
|
||||
index_to_sample_builder = create_mmap_dataset_builder(index_to_sample_fname, sample_idx_dtype)
|
||||
index_to_metric_fname = f"{metric_save_path}/{metric_name}_index_to_metric"
|
||||
index_to_metric_builder = create_mmap_dataset_builder(index_to_metric_fname, metric_value_dtype)
|
||||
for t_idx_reduce in range(num_threads_reduce):
|
||||
chunk_is_fname = f"{metric_save_path}/{metric_name}_index_to_sample_thread{t_idx_reduce}"
|
||||
logger.info(f"Merging file {chunk_is_fname}")
|
||||
index_to_sample_builder.merge_file_(chunk_is_fname)
|
||||
chunk_im_fname = f"{metric_save_path}/{metric_name}_index_to_metric_thread{t_idx_reduce}"
|
||||
logger.info(f"Merging file {chunk_im_fname}")
|
||||
index_to_metric_builder.merge_file_(chunk_im_fname)
|
||||
close_mmap_dataset_builder(index_to_sample_builder, index_to_sample_fname)
|
||||
close_mmap_dataset_builder(index_to_metric_builder, index_to_metric_fname)
|
||||
|
||||
num_sample_per_value = DataAnalyzer.output_index_to_sample_percentile(
|
||||
index_to_sample_fname, index_to_metric_fname, metric_name, metric_save_path, total_num_samples,
|
||||
sample_idx_dtype)
|
||||
self.get_metric_value_percentiles(metric_name, num_sample_per_value, total_num_samples)
|
||||
elif metric_type == 'accumulate_value_over_samples':
|
||||
metric_save_path = f"{save_path}/{metric_name}/"
|
||||
metric_value = None
|
||||
for w_idx in range(num_workers):
|
||||
for t_idx in range(num_threads):
|
||||
w_metric_save_path = f"{metric_save_path}/worker{w_idx}_thread{t_idx}/"
|
||||
w_metric_value_fname = f"{w_metric_save_path}/{metric_name}_metric_value"
|
||||
w_metric_value = MMapIndexedDataset(w_metric_value_fname, skip_warmup=True)
|
||||
if metric_value is None:
|
||||
metric_value = np.copy(w_metric_value[0])
|
||||
else:
|
||||
metric_value += np.copy(w_metric_value[0])
|
||||
value_max = int(max(metric_value))
|
||||
value_min = int(min(metric_value))
|
||||
metric_value_dtype = find_fit_int_dtype(value_min, value_max)
|
||||
metric_value_fname = f"{metric_save_path}/{metric_name}_metric_value"
|
||||
metric_value_builder = create_mmap_dataset_builder(metric_value_fname, metric_value_dtype)
|
||||
metric_value_builder.add_item(torch.tensor(metric_value.astype(np.int64), dtype=torch.long))
|
||||
close_mmap_dataset_builder(metric_value_builder, metric_value_fname)
|
||||
|
||||
@staticmethod
|
||||
def output_index_to_sample_percentile(index_to_sample_fname, index_to_metric_fname, metric_name, metric_save_path,
|
||||
total_num_samples, sample_idx_dtype):
|
||||
""" read index_to_metric and index_to_sample files and write distribution to index_to_sample_percentage_merged """
|
||||
num_sample_per_value = {}
|
||||
index_to_sample = MMapIndexedDataset(index_to_sample_fname, skip_warmup=True)
|
||||
index_to_metric = MMapIndexedDataset(index_to_metric_fname, skip_warmup=True)
|
||||
index_to_sample_merged_fname = f"{metric_save_path}/{metric_name}_index_to_sample_percentile_merged"
|
||||
index_to_sample_merged_builder = create_mmap_dataset_builder(index_to_sample_merged_fname, sample_idx_dtype)
|
||||
for v_idx in range(len(index_to_sample)):
|
||||
if v_idx > 0:
|
||||
assert index_to_metric[v_idx] > index_to_metric[v_idx - 1]
|
||||
num_sample_per_value[index_to_metric[v_idx][0]] = len(index_to_sample[v_idx])
|
||||
assert sum(list(num_sample_per_value.values())) == total_num_samples
|
||||
merge_step = max(1, len(index_to_sample) // 100)
|
||||
for v_idx in range(0, len(index_to_sample), merge_step):
|
||||
merged_samples = np.copy(
|
||||
np.concatenate(index_to_sample[v_idx:min(len(index_to_sample), (v_idx + merge_step))], axis=None))
|
||||
index_to_sample_merged_builder.add_item(torch.tensor(merged_samples.astype(np.int64), dtype=torch.long))
|
||||
logger.info(f"Finished merging index_to_sample {v_idx} to {v_idx+merge_step}.")
|
||||
close_mmap_dataset_builder(index_to_sample_merged_builder, index_to_sample_merged_fname)
|
||||
return num_sample_per_value
|
||||
|
||||
def run_reduce(self):
|
||||
if self.custom_reduce is None:
|
||||
self.merge_map_results(self.dataset, self.metric_names, self.metric_types, self.save_path,
|
||||
self.num_workers, self.num_threads, self.num_threads_reduce)
|
||||
else:
|
||||
self.custom_reduce(self.dataset, self.metric_names, self.metric_types, self.save_path, self.num_workers,
|
||||
self.num_threads, self.num_threads_reduce)
|
||||
|
||||
def run_map_reduce(self, comm_group=None):
|
||||
self.run_map()
|
||||
# wait for the mapping operation, where all nodes outputs their own (partial) result files
|
||||
dist.barrier(group=comm_group)
|
||||
if self.worker_id == 0:
|
||||
self.run_reduce()
|
||||
# wait for the reduce, where rank 0 merges all (partial) files. Dataset can then be used by all nodes.
|
||||
dist.barrier(group=comm_group)
|
||||
|
||||
|
||||
class DistributedDataAnalyzer(object):
|
||||
|
||||
def __init__(
|
||||
self,
|
||||
dataset,
|
||||
num_workers=1,
|
||||
num_threads=1,
|
||||
worker_id=0,
|
||||
batch_size=1,
|
||||
metric_names=[],
|
||||
metric_functions=[],
|
||||
metric_types=[],
|
||||
save_path="./",
|
||||
collate_fn=None,
|
||||
device='cuda',
|
||||
comm_group=None,
|
||||
sample_indices=None,
|
||||
) -> None:
|
||||
self.dataset = dataset
|
||||
self.batch_size = batch_size
|
||||
self.metric_names = metric_names
|
||||
self.metric_functions = metric_functions
|
||||
self.metric_types = metric_types
|
||||
self.save_path = save_path
|
||||
self.collate_fn = collate_fn
|
||||
self.device = device
|
||||
self.sample_indices = sample_indices
|
||||
self.num_threads = num_threads
|
||||
self.worker_id = worker_id
|
||||
|
||||
if not dist.is_initialized():
|
||||
dist.init_distributed()
|
||||
|
||||
# comm_group and worker_id+num_workers are mutually exclusive
|
||||
self.comm_group = comm_group
|
||||
if self.comm_group is None:
|
||||
# self.comm_group = deepspeed.utils.groups._clone_world_group()
|
||||
self.num_workers = num_workers
|
||||
self.worker_id = worker_id
|
||||
else:
|
||||
self.num_workers = self.comm_group.size()
|
||||
self.worker_id = self.comm_group.rank()
|
||||
|
||||
if self.worker_id == 0:
|
||||
logger.info(f"Distributed data analyzer initialized with {self.num_workers} workers.")
|
||||
|
||||
def run_map_helper(self, thread_id=0, metric_queues=None):
|
||||
thread_start_idx, thread_end_idx = self.thread_splits[thread_id][0], self.thread_splits[thread_id][1]
|
||||
worker_dataset = Subset(self.dataset, list(range(thread_start_idx, thread_end_idx)))
|
||||
sampler = BatchSampler(SequentialSampler(worker_dataset), batch_size=self.batch_size, drop_last=False)
|
||||
dataloader = DataLoader(dataset=worker_dataset,
|
||||
batch_sampler=sampler,
|
||||
num_workers=0,
|
||||
collate_fn=self.collate_fn,
|
||||
pin_memory=False)
|
||||
|
||||
# set initial results list
|
||||
metric_results = []
|
||||
for metric_type in self.metric_types:
|
||||
assert metric_type in ['single_value_per_sample', 'accumulate_value_over_samples'], \
|
||||
f"metric_type {metric_type} not implemented."
|
||||
metric_results.append([] if metric_type == 'single_value_per_sample' else None)
|
||||
|
||||
# iterate dataloader and store metric results
|
||||
batch_start_idx = thread_start_idx
|
||||
for data in dataloader:
|
||||
for m_idx in range(len(self.metric_names)):
|
||||
metric_type, metric_function = self.metric_types[m_idx], self.metric_functions[m_idx]
|
||||
metric_values = metric_function(data)
|
||||
assert torch.is_tensor(metric_values) or isinstance(metric_values, np.ndarray), \
|
||||
"metric_function must return a tensor or array"
|
||||
if isinstance(metric_values, np.ndarray):
|
||||
metric_values = torch.from_numpy(metric_values)
|
||||
assert metric_values.dtype in valid_dtypes, \
|
||||
f"metric_function result dtype {metric_values.dtype} not supported. Supported dtypes {valid_dtypes}"
|
||||
|
||||
if metric_type == 'single_value_per_sample':
|
||||
for row in range(metric_values.size()[0]):
|
||||
value = metric_values[row].item()
|
||||
sample_idx = batch_start_idx + row # sample idx following dataset iteration order
|
||||
if isinstance(data, dict) and 'index' in data: # Megatron use case
|
||||
sample_idx = data['index'][row][0].item()
|
||||
elif self.sample_indices is not None: # user defined shuffling of indices
|
||||
sample_idx = self.sample_indices[sample_idx]
|
||||
metric_results[m_idx].append((value, sample_idx))
|
||||
elif metric_type == 'accumulate_value_over_samples':
|
||||
if metric_results[m_idx] is None:
|
||||
metric_results[m_idx] = metric_values
|
||||
else:
|
||||
metric_results[m_idx].add_(metric_values)
|
||||
batch_start_idx += len(data)
|
||||
|
||||
if self.num_threads == 1:
|
||||
return metric_results
|
||||
|
||||
# copy metric_results to the shared queue
|
||||
assert metric_queues
|
||||
for m_idx in range(len(self.metric_names)):
|
||||
results = metric_results[m_idx]
|
||||
if torch.is_tensor(results):
|
||||
results = results.item() if results.dim() == 0 else results.tolist()
|
||||
try:
|
||||
metric_queues[m_idx].put((thread_id, results))
|
||||
except Exception as e:
|
||||
logger.error(f"Error putting metric results to queue: {e}")
|
||||
sys.exit(1)
|
||||
|
||||
def run_map_reduce(self):
|
||||
|
||||
# setup individual dataloaders
|
||||
self.worker_splits, self.thread_splits = split_dataset(self.dataset,
|
||||
self.num_workers,
|
||||
self.worker_id,
|
||||
num_threads=self.num_threads)
|
||||
node_start_idx, node_end_idx = self.worker_splits[self.worker_id]
|
||||
logger.info(f"worker {self.worker_id} working on data subset {node_start_idx} to {node_end_idx}.")
|
||||
|
||||
if self.num_threads in [0, 1, None]:
|
||||
metric_results = self.run_map_helper()
|
||||
metric_results = [torch.tensor(m).to(self.device) for m in metric_results]
|
||||
else:
|
||||
|
||||
# create a shared queue of results per metric to be populated by individual threads
|
||||
with Manager() as manager:
|
||||
metric_queues = [manager.Queue() for _ in self.metric_names]
|
||||
threads = [
|
||||
Process(target=self.run_map_helper, args=(t, metric_queues)) for t in range(self.num_threads)
|
||||
]
|
||||
for thread in threads:
|
||||
thread.start()
|
||||
for thread in threads:
|
||||
thread.join()
|
||||
|
||||
# gather results from shared queues into metric_results
|
||||
metric_results = [None for _ in self.metric_names]
|
||||
for m_idx, (queue, metric_type) in enumerate(zip(metric_queues, self.metric_types)):
|
||||
while not queue.empty():
|
||||
t_idx, t_results = queue.get()
|
||||
t_start_idx, t_end_idx = self.thread_splits[t_idx]
|
||||
if t_start_idx >= t_end_idx: # no results from this thread
|
||||
continue #corner case for small datasets and high thread count
|
||||
t_results = torch.tensor(t_results)
|
||||
if metric_type == 'single_value_per_sample':
|
||||
# add thread results to the metric_results list, ordered by thread idx
|
||||
if metric_results[m_idx] is None: # initialize if needed
|
||||
metric_results[m_idx] = torch.zeros(node_end_idx - node_start_idx,
|
||||
t_results.size(1)).to(self.device)
|
||||
metric_results[m_idx][t_start_idx - node_start_idx:t_end_idx - node_start_idx] = t_results
|
||||
else:
|
||||
if metric_results[m_idx] is None: # initialize if needed
|
||||
metric_results[m_idx] = torch.zeros(t_results.size()).to(self.device)
|
||||
metric_results[m_idx].add_(t_results)
|
||||
|
||||
# compute dtype for sample ids
|
||||
total_num_samples = len(self.dataset)
|
||||
sample_idx_dtype = find_fit_int_dtype(0, total_num_samples - 1)
|
||||
logger.info(f"Total number of data samples: {total_num_samples}.")
|
||||
logger.info(f"Will use {sample_idx_dtype} to store the sample indexes.")
|
||||
|
||||
for m_idx in range(len(self.metric_names)):
|
||||
metric_values, metric_name, metric_type = \
|
||||
metric_results[m_idx], self.metric_names[m_idx], self.metric_types[m_idx]
|
||||
metric_save_path = f"{self.save_path}/{metric_name}/"
|
||||
os.makedirs(metric_save_path, exist_ok=True)
|
||||
|
||||
if metric_type == 'single_value_per_sample':
|
||||
|
||||
# Compute sample and metric value dtypes based on range
|
||||
values, samples = metric_values[:, 0], metric_values[:, 1]
|
||||
value_min, value_max = Dist.min_max(values, self.comm_group)
|
||||
sample_min, sample_max = Dist.min_max(samples, self.comm_group)
|
||||
metric_value_dtype = find_fit_int_dtype(value_min, value_max)
|
||||
sample_value_dtype = find_fit_int_dtype(sample_min, sample_max)
|
||||
|
||||
# sample_to_metric maps sample ids to metric values, as a list of metric values
|
||||
sample_to_metric_fname = f"{metric_save_path}/{metric_name}_sample_to_metric"
|
||||
values = [torch.tensor([x]) for x in metric_values[:, 0]]
|
||||
self.file_write_ordered(values, sample_to_metric_fname, metric_value_dtype)
|
||||
|
||||
# distributed sorting by values, gives an ordered disjoint subset of keys on nodes
|
||||
metric_values = Dist.sample_sort(metric_values, self.comm_group, self.num_workers)
|
||||
metric_to_samples_dict = {}
|
||||
if len(metric_values) > 0:
|
||||
for value, sample in metric_values:
|
||||
if value.item() not in metric_to_samples_dict:
|
||||
metric_to_samples_dict[value.item()] = []
|
||||
metric_to_samples_dict[value.item()].append(sample.item())
|
||||
|
||||
# index_to_metric and index_to_sample serialize a dicitonary from metric to samples
|
||||
# index_to_metric stores a key per row, index_to_sample stores the values per row
|
||||
values = [torch.tensor([x]) for x in metric_to_samples_dict.keys()]
|
||||
samples = [torch.tensor(metric_to_samples_dict[x]) for x in metric_to_samples_dict.keys()]
|
||||
index_to_metric_fname = f"{metric_save_path}/{metric_name}_index_to_metric" #dict keys
|
||||
index_to_sample_fname = f"{metric_save_path}/{metric_name}_index_to_sample" #dict values
|
||||
self.file_write_ordered(values, index_to_metric_fname, metric_value_dtype)
|
||||
self.file_write_ordered(samples, index_to_sample_fname, sample_value_dtype)
|
||||
|
||||
if self.worker_id == 0:
|
||||
DataAnalyzer.output_index_to_sample_percentile(index_to_sample_fname, index_to_metric_fname,
|
||||
metric_name, metric_save_path, total_num_samples,
|
||||
sample_idx_dtype)
|
||||
dist.barrier(self.comm_group)
|
||||
|
||||
elif metric_type == 'accumulate_value_over_samples':
|
||||
metric_value_fname = f"{metric_save_path}/{metric_name}_metric_value"
|
||||
dist.reduce(metric_values, dst=0, op=dist.ReduceOp.SUM, group=self.comm_group)
|
||||
metric_value_dtype = find_fit_int_dtype(metric_values.min(), metric_values.max())
|
||||
|
||||
if self.worker_id == 0:
|
||||
builder = create_mmap_dataset_builder(metric_value_fname, metric_value_dtype)
|
||||
builder.add_item(metric_values.cpu())
|
||||
close_mmap_dataset_builder(builder, metric_value_fname)
|
||||
dist.barrier(self.comm_group)
|
||||
|
||||
def file_write_ordered(self, tensor_list, fname, numpy_dtype):
|
||||
""" MPI_file_write_ordered extended to write a list of tensors, by one rank, iteratively """
|
||||
|
||||
# each node has a list of rows (tensors) to be written to the file.
|
||||
# we will serialize it in order to communicate it in one comm step.
|
||||
|
||||
tkwargs = dict(dtype=torch.int64, device=self.device)
|
||||
|
||||
# 1. gather on rank 0 the number of rows to be sent/recv
|
||||
row_count = torch.tensor([len(tensor_list)], **tkwargs)
|
||||
row_counts = torch.zeros(self.num_workers, **tkwargs)
|
||||
dist.all_gather_into_tensor(row_counts, row_count, group=self.comm_group)
|
||||
assert row_counts[self.worker_id] == row_count == len(tensor_list), "all_gather failed"
|
||||
|
||||
# 2. gather on rank 0 the sizes of the rows to be sent/recv
|
||||
row_len = torch.tensor([len(l) for l in tensor_list], **tkwargs)
|
||||
row_lens = Dist.gather_v(row_len, 0, self.comm_group, self.num_workers, self.worker_id)
|
||||
|
||||
# 4. gather on rank 0 of the total size (sum of all row lengths) to be received
|
||||
size = torch.tensor([sum(row_len).item()], **tkwargs)
|
||||
sizes = torch.zeros(self.num_workers, **tkwargs)
|
||||
dist.all_gather_into_tensor(sizes, size, group=self.comm_group)
|
||||
assert sizes[self.worker_id] == size.item(), "all_gather did not return the same sizes" #sanity check
|
||||
|
||||
# method to deserializes a buffer into rows of different lengths and write them to file
|
||||
def write_buffer_to_file(buff, src, builder):
|
||||
assert self.worker_id == 0, "only rank 0 can write to file"
|
||||
|
||||
# collect all buffers and write them at once
|
||||
buff = buff.cpu().detach().numpy()
|
||||
row_offsets = np.cumsum([0] + row_lens[src].tolist())
|
||||
arr_list = []
|
||||
for i in range(len(row_lens[src])):
|
||||
arr_list.append(buff[row_offsets[i]:row_offsets[i + 1]])
|
||||
builder.add_items(arr_list)
|
||||
|
||||
# 5. rank 0 prepares output folder and file
|
||||
if self.worker_id == 0:
|
||||
os.makedirs(os.path.dirname(fname), exist_ok=True)
|
||||
builder = create_mmap_dataset_builder(fname, numpy_dtype)
|
||||
|
||||
# iterate through ranks that have data to be sent/recv/written
|
||||
for src in [rank for rank, count in enumerate(row_counts) if count > 0]:
|
||||
|
||||
dist.barrier(group=self.comm_group)
|
||||
if self.worker_id == 0 and src == 0: # rank 0's write its own data
|
||||
buffer = torch.cat(tensor_list, dim=0).to(self.device)
|
||||
write_buffer_to_file(buffer, 0, builder)
|
||||
elif self.worker_id == 0 and src > 0: # rank 0 receives other rank's data and writes it
|
||||
buffer = torch.empty(sizes[src].item(), dtype=numpy_dtype, device=self.device)
|
||||
err = dist.recv(buffer, src=src, group=self.comm_group, tag=src)
|
||||
assert err == src and len(buffer) > 0, "recv failed"
|
||||
write_buffer_to_file(buffer, src, builder)
|
||||
elif self.worker_id == src: # current rank sends data to rank 0
|
||||
buffer = torch.cat(tensor_list, dim=0).to(self.device)
|
||||
dist.send(buffer, 0, group=self.comm_group, tag=src)
|
||||
|
||||
# rank 0 closes the file
|
||||
if self.worker_id == 0:
|
||||
close_mmap_dataset_builder(builder, fname) # close file
|
||||
dist.barrier(self.comm_group)
|
||||
|
||||
|
||||
class Dist:
|
||||
""" auxiliary class to perform distributed operations on tensors"""
|
||||
|
||||
@staticmethod
|
||||
def min_max(tensor, comm_group):
|
||||
""" given a distributed tensor, return the min/max values across all ranks"""
|
||||
|
||||
value_min, value_max = tensor.min(), tensor.max()
|
||||
dist.reduce(value_min, 0, op=dist.ReduceOp.MIN, group=comm_group)
|
||||
dist.reduce(value_max, 0, op=dist.ReduceOp.MAX, group=comm_group)
|
||||
return value_min.item(), value_max.item()
|
||||
|
||||
@staticmethod
|
||||
def gather_v(tensor, dst, comm_group, num_workers, worker_id):
|
||||
""" MPI_Gatherv. gather tensors of variable sizes in a single rank """
|
||||
|
||||
# gather the number of rows to be sent/recv
|
||||
size = torch.tensor([len(tensor)], dtype=torch.int64, device=tensor.device)
|
||||
sizes = torch.zeros(num_workers, dtype=torch.int64, device=tensor.device)
|
||||
dist.all_gather_into_tensor(sizes, size, group=comm_group)
|
||||
assert sizes[worker_id] == size, "all_gather failed"
|
||||
|
||||
# all_gather requires all tensors to be of same size so we need to pad them
|
||||
max_size = max(sizes).item()
|
||||
buffer = torch.empty(max_size, dtype=tensor.dtype, device=tensor.device)
|
||||
buffer[0:size] = tensor.data
|
||||
buffer_list = None
|
||||
if worker_id == 0: # create padded recv buffers
|
||||
buffer_list = [torch.empty(max_size, dtype=tensor.dtype, device=tensor.device) for _ in range(num_workers)]
|
||||
dist.gather(buffer, buffer_list, dst=dst, group=comm_group)
|
||||
|
||||
# revert padding and return value
|
||||
if worker_id == 0:
|
||||
buffer_list = [r[:s.item()] for r, s in zip(buffer_list, sizes)]
|
||||
return buffer_list
|
||||
|
||||
@staticmethod
|
||||
def sample_sort(tensor, comm_group, num_workers, n_samples=100):
|
||||
""" perform a distributed random sort of a tensor, and returns the sorted partial tensor"""
|
||||
device, dims = tensor.device, tensor.size()[1]
|
||||
|
||||
# 1 - sort rows by first column, then second column, then third, etc...
|
||||
tensor = torch.tensor(sorted(tensor.tolist()), dtype=tensor.dtype, device=tensor.device)
|
||||
|
||||
# 2 - collect few samples per rank
|
||||
idx = torch.round(torch.linspace(0, len(tensor) - 1, n_samples)).to(int)
|
||||
samples = tensor[idx][:, 0].contiguous().to(device) #only first column, all but last row
|
||||
|
||||
# 2 - Allgather samples
|
||||
all_samples = [torch.zeros(n_samples, dtype=samples.dtype, device=device) for _ in range(num_workers)]
|
||||
dist.all_gather(all_samples, samples, group=comm_group)
|
||||
all_samples = torch.cat(all_samples, dim=0).to(device)
|
||||
|
||||
# 3 - Sort all samples and collect the ranges of each rank as equidistant
|
||||
all_samples = all_samples.sort()[0]
|
||||
idx = torch.round(torch.linspace(0, len(all_samples) - 1, num_workers + 1)).to(int)
|
||||
ranges = all_samples[idx] # range of each rank r as ranges[r] <= x < ranges[r+1]
|
||||
ranges[-1] += 1 # increase upper limit of last rank so that x < ranges[r+1].
|
||||
|
||||
# 4 - collect elements to send to each rank, based on the rank ranges
|
||||
send = []
|
||||
for rank in range(num_workers):
|
||||
mask = (tensor[:, 0] >= ranges[rank]) & (tensor[:, 0] < ranges[rank + 1])
|
||||
send.append(tensor[mask])
|
||||
|
||||
# 5. all to all to communicate the sizes to be sent/recv
|
||||
send_count = [torch.tensor([len(s) * dims], dtype=torch.int64, device=device) for s in send]
|
||||
recv_count = list(torch.empty([num_workers], dtype=torch.int64, device=device).chunk(num_workers))
|
||||
dist.all_to_all(recv_count, send_count, group=comm_group)
|
||||
|
||||
# 6. all-to-all-v to communicate the elements to be sent/recv as a single tensor
|
||||
send = torch.cat(send, dim=0).flatten().to(device)
|
||||
recv = torch.zeros(sum(recv_count), dtype=send.dtype).to(device)
|
||||
send_count = [s.item() for s in send_count] # convert to list of ints
|
||||
recv_count = [r.item() for r in recv_count]
|
||||
dist.all_to_all_single(recv, send, recv_count, send_count, group=comm_group)
|
||||
del send
|
||||
|
||||
# 7. the received tensor is the 1D disjoint subset of the distributed tensor.
|
||||
# We will recover the original dimensionality and sort it by columns again.
|
||||
recv = recv.view(-1, dims)
|
||||
recv = torch.tensor(sorted(recv.tolist()), dtype=recv.dtype, device=recv.device)
|
||||
return recv
|
||||
|
||||
|
||||
def test_compare_both_data_analyzers(dataset):
|
||||
""" given a dataset, compare file and memory based data analyser"""
|
||||
|
||||
id = lambda t: t.to(torch.int64) # identity
|
||||
batch_sum = lambda t: id(t).sum() #sum batch
|
||||
num_threads = 4
|
||||
kwargs = dict(
|
||||
dataset=dataset,
|
||||
batch_size=2**10,
|
||||
worker_id=int(os.environ['RANK']),
|
||||
num_workers=int(os.environ['WORLD_SIZE']),
|
||||
metric_names=["mod", "batch_sum"],
|
||||
metric_functions=[id, batch_sum],
|
||||
metric_types=['single_value_per_sample', 'accumulate_value_over_samples'],
|
||||
num_threads=num_threads,
|
||||
)
|
||||
|
||||
dda = DistributedDataAnalyzer(
|
||||
save_path="./output_dist",
|
||||
device=f"cuda:{int(os.environ['LOCAL_RANK'])}",
|
||||
**kwargs,
|
||||
)
|
||||
start_time = time.time()
|
||||
dda.run_map_reduce()
|
||||
if dda.worker_id == 0:
|
||||
print("DistributedDataAnalyzer runtime: %s seconds " % (time.time() - start_time))
|
||||
|
||||
da = DataAnalyzer(num_threads_reduce=num_threads,
|
||||
save_path="./output_disk",
|
||||
metric_dtypes=[torch.int64, torch.int64],
|
||||
**kwargs)
|
||||
start_time = time.time()
|
||||
da.run_map_reduce()
|
||||
if da.worker_id == 0:
|
||||
print("DataAnalyzer runtime: %s seconds " % (time.time() - start_time))
|
||||
|
||||
output_paths = [
|
||||
"batch_sum/batch_sum_metric_value.bin", "batch_sum/batch_sum_metric_value.idx", \
|
||||
"mod/mod_index_to_metric.bin", "mod/mod_index_to_metric.idx", \
|
||||
"mod/mod_index_to_sample.bin", "mod/mod_index_to_sample.idx", \
|
||||
"mod/mod_index_to_sample_percentile_merged.bin", "mod/mod_index_to_sample_percentile_merged.idx", \
|
||||
"mod/mod_sample_to_metric.bin", "mod/mod_sample_to_metric.idx"
|
||||
]
|
||||
|
||||
if dda.worker_id == 0:
|
||||
for path in output_paths:
|
||||
with open(os.path.join(da.save_path, path), 'rb') as f1, \
|
||||
open(os.path.join(dda.save_path, path), 'rb') as f2:
|
||||
# if files have suffix .bin, they should be identical
|
||||
if path.endswith(".bin"):
|
||||
assert f1.read() == f2.read(), f"files {path} are not identical."
|
||||
elif f1.read() != f2.read():
|
||||
print(f"files {path} are not identical.")
|
||||
dist.barrier()
|
||||
dist.destroy_process_group()
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
|
||||
class TestDataset(torch.utils.data.Dataset):
|
||||
|
||||
def __init__(self, size=10_000_000):
|
||||
self.values = [(x + 7) % 10_000 for x in range(size)]
|
||||
self.size = size
|
||||
|
||||
__len__ = lambda self: self.size
|
||||
__getitem__ = lambda self, idx: self.values[idx]
|
||||
|
||||
test_compare_both_data_analyzers(TestDataset())
|
||||
@@ -0,0 +1,349 @@
|
||||
# Copyright (c) Microsoft Corporation.
|
||||
# SPDX-License-Identifier: Apache-2.0
|
||||
|
||||
# DeepSpeed Team
|
||||
"""
|
||||
coding=utf-8
|
||||
Copyright (c) 2020, NVIDIA CORPORATION. All rights reserved.
|
||||
|
||||
Licensed under the Apache License, Version 2.0 (the "License");
|
||||
you may not use this file except in compliance with the License.
|
||||
You may obtain a copy of the License at
|
||||
|
||||
http://www.apache.org/licenses/LICENSE-2.0
|
||||
|
||||
Unless required by applicable law or agreed to in writing, software
|
||||
distributed under the License is distributed on an "AS IS" BASIS,
|
||||
WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
|
||||
See the License for the specific language governing permissions and
|
||||
limitations under the License.
|
||||
Part of this code was adopted from https://github.com/NVIDIA/Megatron-LM/blob/main/megatron/data/data_samplers.py
|
||||
"""
|
||||
|
||||
import torch
|
||||
import os
|
||||
import numpy as np
|
||||
|
||||
import deepspeed.comm as dist
|
||||
from deepspeed.utils import logger
|
||||
from deepspeed.accelerator import get_accelerator
|
||||
from ..constants import *
|
||||
from ..curriculum_scheduler import CurriculumScheduler
|
||||
from .indexed_dataset import MMapIndexedDataset
|
||||
from .utils import create_mmap_dataset_builder, close_mmap_dataset_builder, find_fit_int_dtype
|
||||
|
||||
|
||||
class DeepSpeedDataSampler(object):
|
||||
|
||||
def __init__(self,
|
||||
data_efficiency_config,
|
||||
one_epoch_total_samples,
|
||||
micro_batch_size,
|
||||
data_parallel_rank,
|
||||
data_parallel_size,
|
||||
data_parallel_group,
|
||||
gradient_accumulation_steps,
|
||||
global_rank,
|
||||
drop_last=True):
|
||||
# Keep a copy of input params for later use.
|
||||
self.data_efficiency_config = data_efficiency_config
|
||||
self.one_epoch_total_samples = one_epoch_total_samples
|
||||
self.index_dtype = find_fit_int_dtype(0, one_epoch_total_samples)
|
||||
self.total_samples = one_epoch_total_samples * self.data_efficiency_config[DATA_SAMPLING][
|
||||
DATA_SAMPLING_NUM_EPOCHS]
|
||||
self.micro_batch_size = micro_batch_size
|
||||
self.data_parallel_rank = data_parallel_rank
|
||||
self.data_parallel_group = data_parallel_group
|
||||
self.micro_batch_times_data_parallel_size = \
|
||||
self.micro_batch_size * data_parallel_size
|
||||
self.gradient_accumulation_steps = gradient_accumulation_steps
|
||||
self.global_batch_size = self.micro_batch_times_data_parallel_size * \
|
||||
self.gradient_accumulation_steps
|
||||
self.global_rank = global_rank
|
||||
self.drop_last = drop_last
|
||||
self.np_rng = np.random.default_rng(self.data_efficiency_config[DATA_EFFICIENCY_SEED])
|
||||
self.state = {}
|
||||
self.batch = []
|
||||
self.consumed_samples = 0
|
||||
if self.data_efficiency_config[DATA_SAMPLING][CURRICULUM_LEARNING][CURRICULUM_LEARNING_ENABLED]:
|
||||
self.curriculum_step = 0
|
||||
self.current_difficulties = {}
|
||||
self.data_cluster_paths = []
|
||||
self.data_cluster_current_position = []
|
||||
self.curriculum_schedulers = {}
|
||||
self.curriculum_index_to_sample = {}
|
||||
self.curriculum_index_to_metric = {}
|
||||
self.difficulty_type = {}
|
||||
self.clustering_type = {}
|
||||
self.data_1epoch_size = None
|
||||
if self.global_rank == 0:
|
||||
self.data_clusters = []
|
||||
self.data_cluster_sizes = []
|
||||
cluster_path = self.data_efficiency_config[DATA_SAMPLING][CURRICULUM_LEARNING][
|
||||
CURRICULUM_LEARNING_CLUSTER_PATH]
|
||||
if not os.path.exists(cluster_path):
|
||||
os.makedirs(cluster_path)
|
||||
for metric in self.data_efficiency_config[DATA_SAMPLING][CURRICULUM_LEARNING][CURRICULUM_LEARNING_METRICS]:
|
||||
self.curriculum_schedulers[metric] = CurriculumScheduler(
|
||||
data_efficiency_config[DATA_SAMPLING][CURRICULUM_LEARNING][CURRICULUM_LEARNING_METRICS][metric])
|
||||
self.difficulty_type[metric] = data_efficiency_config[DATA_SAMPLING][CURRICULUM_LEARNING][
|
||||
CURRICULUM_LEARNING_METRICS][metric][CURRICULUM_LEARNING_DIFFICULTY_TYPE]
|
||||
self.clustering_type[metric] = data_efficiency_config[DATA_SAMPLING][CURRICULUM_LEARNING][
|
||||
CURRICULUM_LEARNING_METRICS][metric][CURRICULUM_LEARNING_CLUSTERING_TYPE]
|
||||
if self.global_rank == 0:
|
||||
if self.clustering_type[metric] != CURRICULUM_LEARNING_SINGLE_CLUSTER:
|
||||
self.curriculum_index_to_sample[metric] = MMapIndexedDataset(
|
||||
data_efficiency_config[DATA_SAMPLING][CURRICULUM_LEARNING][CURRICULUM_LEARNING_METRICS]
|
||||
[metric][CURRICULUM_LEARNING_SAMPLE_PATH],
|
||||
skip_warmup=True)
|
||||
if self.difficulty_type[metric] == CURRICULUM_LEARNING_VALUE_BASED:
|
||||
self.curriculum_index_to_metric[metric] = MMapIndexedDataset(
|
||||
data_efficiency_config[DATA_SAMPLING][CURRICULUM_LEARNING][CURRICULUM_LEARNING_METRICS]
|
||||
[metric][CURRICULUM_LEARNING_METRIC_PATH],
|
||||
skip_warmup=True)
|
||||
|
||||
# Sanity checks.
|
||||
assert self.total_samples > 0, \
|
||||
'no sample to consume: {}'.format(self.total_samples)
|
||||
assert self.micro_batch_size > 0
|
||||
assert data_parallel_size > 0
|
||||
assert self.data_parallel_rank < data_parallel_size, \
|
||||
'data_parallel_rank should be smaller than data size: {}, ' \
|
||||
'{}'.format(self.data_parallel_rank, data_parallel_size)
|
||||
|
||||
def __len__(self):
|
||||
return self.total_samples
|
||||
|
||||
def set_custom_curriculum_learning_schedule(self, schedule_func_dict):
|
||||
for metric in self.curriculum_schedulers:
|
||||
if metric in schedule_func_dict:
|
||||
self.curriculum_schedulers[metric].set_custom_get_difficulty(schedule_func_dict[metric])
|
||||
|
||||
def get_start_end_idx(self, batch_len=None):
|
||||
"""
|
||||
given the length of a minibatch (defaults to micro-batch size * data_parallel_size),
|
||||
return the start and end indices of the current data parallel rank
|
||||
"""
|
||||
batch_len = batch_len or self.micro_batch_times_data_parallel_size
|
||||
start_idx_fn = lambda r: round(r * batch_len / self.data_parallel_group.size())
|
||||
start_idx = start_idx_fn(self.data_parallel_rank)
|
||||
end_idx = start_idx_fn(self.data_parallel_rank + 1)
|
||||
return start_idx, end_idx
|
||||
|
||||
def get_sample_based_on_metric_value(self, metric, value_start, value_end):
|
||||
new_samples = None
|
||||
for row in range(len(self.curriculum_index_to_sample[metric])):
|
||||
if self.curriculum_index_to_metric[metric][row] <= value_end and self.curriculum_index_to_metric[metric][
|
||||
row] > value_start:
|
||||
row_samples = np.copy(self.curriculum_index_to_sample[metric][row])
|
||||
new_samples = row_samples if new_samples is None else np.concatenate(
|
||||
(new_samples, row_samples), axis=None)
|
||||
return new_samples
|
||||
|
||||
def get_sample_based_on_metric_percentile(self, metric, percentile_start, percentile_end):
|
||||
new_samples = None
|
||||
if self.data_1epoch_size is None:
|
||||
self.data_1epoch_size = sum(len(x) for x in self.curriculum_index_to_sample[metric])
|
||||
max_percentile = self.data_efficiency_config[DATA_SAMPLING][CURRICULUM_LEARNING][CURRICULUM_LEARNING_METRICS][
|
||||
metric][CURRICULUM_LEARNING_MAX_DIFFICULTY]
|
||||
sample_per_percentile = self.data_1epoch_size // max_percentile
|
||||
start_count = sample_per_percentile * percentile_start
|
||||
end_count = sample_per_percentile * percentile_end
|
||||
if percentile_end == max_percentile:
|
||||
end_count = self.data_1epoch_size
|
||||
current_count = 0
|
||||
for row in range(len(self.curriculum_index_to_sample[metric])):
|
||||
row_size = len(self.curriculum_index_to_sample[metric][row])
|
||||
if current_count + row_size > start_count:
|
||||
row_start = max(0, start_count - current_count)
|
||||
if current_count + row_size <= end_count:
|
||||
row_end = row_size
|
||||
else:
|
||||
row_end = end_count - current_count
|
||||
row_samples = np.copy(self.curriculum_index_to_sample[metric][row][row_start:row_end])
|
||||
new_samples = row_samples if new_samples is None else np.concatenate(
|
||||
(new_samples, row_samples), axis=None)
|
||||
current_count += row_size
|
||||
if current_count >= end_count:
|
||||
break
|
||||
return new_samples
|
||||
|
||||
def get_new_cluster(self, previous_difficulties):
|
||||
cluster_fname = CURRICULUM_LEARNING_CLUSTER_PREFIX
|
||||
for metric in self.curriculum_schedulers:
|
||||
cluster_fname = f"{cluster_fname}_{metric}{self.current_difficulties[metric]}"
|
||||
cluster_path = self.data_efficiency_config[DATA_SAMPLING][CURRICULUM_LEARNING][
|
||||
CURRICULUM_LEARNING_CLUSTER_PATH]
|
||||
cluster_path = f"{cluster_path}/{cluster_fname}"
|
||||
if self.global_rank == 0:
|
||||
new_cluster = None
|
||||
need_clustering = 0
|
||||
for metric in self.clustering_type:
|
||||
if self.clustering_type[metric] != CURRICULUM_LEARNING_SINGLE_CLUSTER:
|
||||
need_clustering += 1
|
||||
if need_clustering > 1:
|
||||
for metric in self.curriculum_schedulers:
|
||||
if self.clustering_type[metric] == CURRICULUM_LEARNING_SINGLE_CLUSTER:
|
||||
metric_cluster = np.arange(start=0,
|
||||
stop=self.one_epoch_total_samples,
|
||||
step=1,
|
||||
dtype=self.index_dtype)
|
||||
else:
|
||||
if self.difficulty_type[metric] == CURRICULUM_LEARNING_VALUE_BASED:
|
||||
metric_cluster = self.get_sample_based_on_metric_value(metric, float('-inf'),
|
||||
self.current_difficulties[metric])
|
||||
elif self.difficulty_type[metric] == CURRICULUM_LEARNING_PERCENTILE_BASED:
|
||||
metric_cluster = self.get_sample_based_on_metric_percentile(
|
||||
metric, 0, self.current_difficulties[metric])
|
||||
new_cluster = metric_cluster if new_cluster is None else \
|
||||
np.intersect1d(new_cluster, metric_cluster, assume_unique=True)
|
||||
for cluster in self.data_clusters:
|
||||
new_cluster = np.setdiff1d(new_cluster, cluster[0], assume_unique=True)
|
||||
else:
|
||||
if len(self.data_clusters) == 0:
|
||||
new_cluster = np.arange(start=0, stop=self.one_epoch_total_samples, step=1, dtype=self.index_dtype)
|
||||
for metric in self.curriculum_schedulers:
|
||||
if self.clustering_type[metric] != CURRICULUM_LEARNING_SINGLE_CLUSTER:
|
||||
if self.difficulty_type[metric] == CURRICULUM_LEARNING_VALUE_BASED:
|
||||
new_cluster = self.get_sample_based_on_metric_value(metric, previous_difficulties[metric],
|
||||
self.current_difficulties[metric])
|
||||
elif self.difficulty_type[metric] == CURRICULUM_LEARNING_PERCENTILE_BASED:
|
||||
new_cluster = self.get_sample_based_on_metric_percentile(
|
||||
metric, previous_difficulties[metric], self.current_difficulties[metric])
|
||||
if new_cluster is not None and len(new_cluster) > 0:
|
||||
logger.info(
|
||||
f"new data cluster (previous_difficulties {previous_difficulties}, current_difficulties {self.current_difficulties}) with size {len(new_cluster)} generated."
|
||||
)
|
||||
self.np_rng.shuffle(new_cluster)
|
||||
cluster_builder = create_mmap_dataset_builder(cluster_path, self.index_dtype)
|
||||
cluster_builder.add_item_numpy(new_cluster)
|
||||
close_mmap_dataset_builder(cluster_builder, cluster_path)
|
||||
self.data_clusters.append(MMapIndexedDataset(cluster_path, skip_warmup=True))
|
||||
self.data_cluster_sizes.append(len(self.data_clusters[-1][0]))
|
||||
else:
|
||||
logger.info(
|
||||
f"new data cluster (previous_difficulties {previous_difficulties}, current_difficulties {self.current_difficulties}) has no matched data thus skipped."
|
||||
)
|
||||
dist.barrier(group=self.data_parallel_group)
|
||||
if os.path.isfile(f"{cluster_path}.bin"):
|
||||
self.data_cluster_paths.append(cluster_fname)
|
||||
self.data_cluster_current_position.append(0)
|
||||
|
||||
def sample_from_clusters(self):
|
||||
num_clusters = len(self.data_clusters)
|
||||
weight_sum = sum(self.data_cluster_sizes)
|
||||
weights = [x / weight_sum for x in self.data_cluster_sizes]
|
||||
samples = self.np_rng.choice(num_clusters, self.global_batch_size, replace=True, p=weights)
|
||||
samples = np.bincount(samples, minlength=num_clusters)
|
||||
return samples
|
||||
|
||||
def reshuffle_clusters(self, cidx):
|
||||
cluster_fname = self.data_cluster_paths[cidx]
|
||||
cluster_path = self.data_efficiency_config[DATA_SAMPLING][CURRICULUM_LEARNING][
|
||||
CURRICULUM_LEARNING_CLUSTER_PATH]
|
||||
cluster_path = f"{cluster_path}/{cluster_fname}"
|
||||
cluster = np.copy(self.data_clusters[cidx][0])
|
||||
self.np_rng.shuffle(cluster)
|
||||
cluster_builder = create_mmap_dataset_builder(cluster_path, self.index_dtype)
|
||||
cluster_builder.add_item_numpy(cluster)
|
||||
close_mmap_dataset_builder(cluster_builder, cluster_path)
|
||||
self.data_clusters[cidx] = MMapIndexedDataset(cluster_path, skip_warmup=True)
|
||||
|
||||
def get_sample_from_cluster(self, cidx, num_samples):
|
||||
start_idx = self.data_cluster_current_position[cidx]
|
||||
samples = list(np.copy(self.data_clusters[cidx][0][start_idx:(start_idx + num_samples)]))
|
||||
self.data_cluster_current_position[cidx] += num_samples
|
||||
if len(samples) < num_samples:
|
||||
num_samples_remained = num_samples - len(samples)
|
||||
logger.info(f"reshuffling cluster {cidx}.")
|
||||
self.reshuffle_clusters(cidx)
|
||||
samples += list(np.copy(self.data_clusters[cidx][0][:num_samples_remained]))
|
||||
self.data_cluster_current_position[cidx] = num_samples_remained
|
||||
return samples
|
||||
|
||||
def get_next_global_batch(self):
|
||||
if self.data_efficiency_config[DATA_SAMPLING][CURRICULUM_LEARNING][CURRICULUM_LEARNING_ENABLED]:
|
||||
self.curriculum_step += 1
|
||||
new_cluster = False
|
||||
previous_difficulties = {}
|
||||
for metric in self.curriculum_schedulers:
|
||||
next_difficulty = self.curriculum_schedulers[metric].update_difficulty(self.curriculum_step)
|
||||
if metric not in self.current_difficulties or \
|
||||
next_difficulty != self.current_difficulties[metric]:
|
||||
new_cluster = True
|
||||
if metric in self.current_difficulties:
|
||||
previous_difficulties[metric] = self.current_difficulties[metric]
|
||||
else:
|
||||
if self.difficulty_type[metric] == CURRICULUM_LEARNING_VALUE_BASED:
|
||||
previous_difficulties[metric] = float('-inf')
|
||||
elif self.difficulty_type[metric] == CURRICULUM_LEARNING_PERCENTILE_BASED:
|
||||
previous_difficulties[metric] = 0
|
||||
self.current_difficulties[metric] = next_difficulty
|
||||
if new_cluster:
|
||||
self.get_new_cluster(previous_difficulties)
|
||||
if self.global_rank == 0:
|
||||
samples_per_cluster = self.sample_from_clusters()
|
||||
batch = []
|
||||
for cidx in range(len(samples_per_cluster)):
|
||||
batch += self.get_sample_from_cluster(cidx, samples_per_cluster[cidx])
|
||||
self.np_rng.shuffle(batch)
|
||||
|
||||
# broadcast tensor must have same shape across participants. So we fill batch with -1s when not full
|
||||
assert len(batch) <= self.global_batch_size
|
||||
batch += [-1] * (self.global_batch_size - len(batch))
|
||||
batch = torch.tensor(batch, device=get_accelerator().current_device_name(), dtype=torch.long).view(-1)
|
||||
else:
|
||||
batch = torch.empty(self.global_batch_size,
|
||||
device=get_accelerator().current_device_name(),
|
||||
dtype=torch.long)
|
||||
dist.broadcast(batch, 0, group=self.data_parallel_group)
|
||||
batch = batch[batch != -1] # remove trailing -1s used to fill incomplete batch tensor
|
||||
self.batch = batch.tolist()
|
||||
|
||||
def __iter__(self):
|
||||
while self.consumed_samples <= self.total_samples:
|
||||
if len(self.batch) == 0:
|
||||
self.get_next_global_batch()
|
||||
current_batch = self.batch[:self.micro_batch_times_data_parallel_size]
|
||||
self.batch = self.batch[self.micro_batch_times_data_parallel_size:]
|
||||
if len(current_batch) == self.micro_batch_times_data_parallel_size or \
|
||||
(len(current_batch) > 0 and not self.drop_last):
|
||||
start_idx, end_idx = self.get_start_end_idx(len(current_batch))
|
||||
yield current_batch[start_idx:end_idx]
|
||||
self.consumed_samples += len(current_batch)
|
||||
current_batch = []
|
||||
|
||||
def state_dict(self):
|
||||
return {
|
||||
CURRICULUM_LEARNING_BATCH: self.batch,
|
||||
CURRICULUM_LEARNING_CONSUMED_SAMPLES: self.consumed_samples,
|
||||
CURRICULUM_LEARNING_STEP: self.curriculum_step,
|
||||
CURRICULUM_LEARNING_CURRENT_DIFFICULTIES: self.current_difficulties,
|
||||
CURRICULUM_LEARNING_DATA_CLUSTER_PATHS: self.data_cluster_paths,
|
||||
CURRICULUM_LEARNING_DATA_CLUSTER_CURRENT_POSITION: self.data_cluster_current_position,
|
||||
CURRICULUM_LEARNING_NP_RNG_STATE: np.random.get_state()
|
||||
}
|
||||
|
||||
def load_state_dict(self, state_dict):
|
||||
self.batch = state_dict[CURRICULUM_LEARNING_BATCH]
|
||||
self.consumed_samples = state_dict[CURRICULUM_LEARNING_CONSUMED_SAMPLES]
|
||||
self.curriculum_step = state_dict[CURRICULUM_LEARNING_STEP]
|
||||
self.current_difficulties = state_dict[CURRICULUM_LEARNING_CURRENT_DIFFICULTIES]
|
||||
self.data_cluster_paths = state_dict[CURRICULUM_LEARNING_DATA_CLUSTER_PATHS]
|
||||
self.data_cluster_current_position = state_dict[CURRICULUM_LEARNING_DATA_CLUSTER_CURRENT_POSITION]
|
||||
np.random.set_state(state_dict[CURRICULUM_LEARNING_NP_RNG_STATE])
|
||||
cluster_root_path = self.data_efficiency_config[DATA_SAMPLING][CURRICULUM_LEARNING][
|
||||
CURRICULUM_LEARNING_CLUSTER_PATH]
|
||||
# Backward compatibility: previously data_cluster_paths were stored as
|
||||
# absolute paths. Now we changed it to just the file name so that even
|
||||
# if user moved the cluster files, the checkpoint loading still works
|
||||
# as long as user set the correct new CURRICULUM_LEARNING_CLUSTER_PATH
|
||||
# in deepspeed json config.
|
||||
for idx in range(len(self.data_cluster_paths)):
|
||||
if '/' in self.data_cluster_paths[idx]:
|
||||
self.data_cluster_paths[idx] = self.data_cluster_paths[idx].split('/')[-1]
|
||||
if self.global_rank == 0:
|
||||
for cluster_fname in self.data_cluster_paths:
|
||||
cluster_path = f"{cluster_root_path}/{cluster_fname}"
|
||||
self.data_clusters.append(MMapIndexedDataset(cluster_path, skip_warmup=True))
|
||||
self.data_cluster_sizes.append(len(self.data_clusters[-1][0]))
|
||||
@@ -0,0 +1,627 @@
|
||||
# Copyright (c) Microsoft Corporation.
|
||||
# SPDX-License-Identifier: Apache-2.0
|
||||
|
||||
# DeepSpeed Team
|
||||
"""
|
||||
Part of this code was adopted from https://github.com/NVIDIA/Megatron-LM/blob/main/megatron/data/indexed_dataset.py
|
||||
"""
|
||||
|
||||
# Copyright (c) Facebook, Inc. and its affiliates.
|
||||
#
|
||||
# This source code is licensed under the MIT license found in the
|
||||
# LICENSE file in the root directory of this source tree.
|
||||
|
||||
# copied from fairseq/fairseq/data/indexed_dataset.py
|
||||
# Removed IndexedRawTextDataset since it relied on Fairseq dictionary
|
||||
# other slight modifications to remove fairseq dependencies
|
||||
# Added document index to index file and made it accessible.
|
||||
# An empty sentence no longer separates documents.
|
||||
|
||||
# Some of the fixes/improvements are adopted from
|
||||
# https://github.com/bigscience-workshop/Megatron-DeepSpeed/blob/main/megatron/data/indexed_dataset.py
|
||||
|
||||
from functools import lru_cache
|
||||
import os
|
||||
import shutil
|
||||
import struct
|
||||
from itertools import accumulate
|
||||
|
||||
import numpy as np
|
||||
import torch
|
||||
|
||||
|
||||
def __best_fitting_dtype(vocab_size=None):
|
||||
if vocab_size is not None and vocab_size < 65500:
|
||||
return np.uint16
|
||||
else:
|
||||
return np.int32
|
||||
|
||||
|
||||
def get_available_dataset_impl():
|
||||
return ['lazy', 'cached', 'mmap']
|
||||
|
||||
|
||||
def infer_dataset_impl(path):
|
||||
if IndexedDataset.exists(path):
|
||||
with open(index_file_path(path), 'rb') as f:
|
||||
magic = f.read(8)
|
||||
if magic == IndexedDataset._HDR_MAGIC:
|
||||
return 'cached'
|
||||
elif magic == MMapIndexedDataset.Index._HDR_MAGIC[:8]:
|
||||
return 'mmap'
|
||||
else:
|
||||
return None
|
||||
else:
|
||||
print(f"Dataset does not exist: {path}")
|
||||
print("Path should be a basename that both .idx and .bin can be appended to get full filenames.")
|
||||
return None
|
||||
|
||||
|
||||
def make_builder(out_file, impl, vocab_size=None):
|
||||
if impl == 'mmap':
|
||||
return MMapIndexedDatasetBuilder(out_file, dtype=__best_fitting_dtype(vocab_size))
|
||||
else:
|
||||
return IndexedDatasetBuilder(out_file)
|
||||
|
||||
|
||||
def make_dataset(path, impl, skip_warmup=False):
|
||||
if not IndexedDataset.exists(path):
|
||||
print(f"Dataset does not exist: {path}")
|
||||
print("Path should be a basename that both .idx and .bin can be appended to get full filenames.")
|
||||
return None
|
||||
if impl == 'infer':
|
||||
impl = infer_dataset_impl(path)
|
||||
if impl == 'lazy' and IndexedDataset.exists(path):
|
||||
return IndexedDataset(path)
|
||||
elif impl == 'cached' and IndexedDataset.exists(path):
|
||||
return IndexedCachedDataset(path)
|
||||
elif impl == 'mmap' and MMapIndexedDataset.exists(path):
|
||||
return MMapIndexedDataset(path, skip_warmup)
|
||||
print(f"Unknown dataset implementation: {impl}")
|
||||
return None
|
||||
|
||||
|
||||
def dataset_exists(path, impl):
|
||||
if impl == 'mmap':
|
||||
return MMapIndexedDataset.exists(path)
|
||||
else:
|
||||
return IndexedDataset.exists(path)
|
||||
|
||||
|
||||
def read_longs(f, n):
|
||||
a = np.empty(n, dtype=np.int64)
|
||||
f.readinto(a)
|
||||
return a
|
||||
|
||||
|
||||
def write_longs(f, a):
|
||||
f.write(np.array(a, dtype=np.int64))
|
||||
|
||||
|
||||
# valid metric_dtypes as numpy and torch types
|
||||
dtypes = {
|
||||
1: (np.uint8, torch.uint8),
|
||||
2: (np.int8, torch.int8),
|
||||
3: (np.int16, torch.int16),
|
||||
4: (np.int32, torch.int32),
|
||||
5: (np.int64, torch.int64),
|
||||
6: (np.uint16, None),
|
||||
7: (np.uint32, None),
|
||||
8: (np.uint64, None),
|
||||
}
|
||||
|
||||
valid_dtypes = set([dt[0] for dt in dtypes.values()] + [dt[1] for dt in dtypes.values() if dt[1] is not None])
|
||||
|
||||
|
||||
def code(dtype):
|
||||
for c, (np_dt, torch_dt) in dtypes.items():
|
||||
if dtype in [np_dt, torch_dt]:
|
||||
return c
|
||||
raise ValueError(f"{dtype} not supported. Supported types: {valid_dtypes}")
|
||||
|
||||
|
||||
def index_file_path(prefix_path):
|
||||
return prefix_path + '.idx'
|
||||
|
||||
|
||||
def data_file_path(prefix_path):
|
||||
return prefix_path + '.bin'
|
||||
|
||||
|
||||
def create_doc_idx(sizes):
|
||||
doc_idx = [0]
|
||||
for i, s in enumerate(sizes):
|
||||
if s == 0:
|
||||
doc_idx.append(i + 1)
|
||||
return doc_idx
|
||||
|
||||
|
||||
class IndexedDataset(torch.utils.data.Dataset):
|
||||
"""Loader for IndexedDataset"""
|
||||
_HDR_MAGIC = b'TNTIDX\x00\x00'
|
||||
|
||||
def __init__(self, path):
|
||||
super().__init__()
|
||||
self.path = path
|
||||
self.data_file = None
|
||||
self.read_index(path)
|
||||
|
||||
def read_index(self, path):
|
||||
with open(index_file_path(path), 'rb') as f:
|
||||
magic = f.read(8)
|
||||
assert magic == self._HDR_MAGIC, ('Index file doesn\'t match expected format. '
|
||||
'Make sure that --dataset-impl is configured properly.')
|
||||
version = f.read(8)
|
||||
assert struct.unpack('<Q', version) == (1, )
|
||||
code, self.element_size = struct.unpack('<QQ', f.read(16))
|
||||
self.dtype = dtypes[code][0] #numpy type
|
||||
self._len, self.s = struct.unpack('<QQ', f.read(16))
|
||||
self.doc_count = struct.unpack('<Q', f.read(8))
|
||||
self.dim_offsets = read_longs(f, self._len + 1)
|
||||
self.data_offsets = read_longs(f, self._len + 1)
|
||||
self.sizes = read_longs(f, self.s)
|
||||
self.doc_idx = read_longs(f, self.doc_count)
|
||||
|
||||
def read_data(self, path):
|
||||
self.data_file = open(data_file_path(path), 'rb', buffering=0)
|
||||
|
||||
def check_index(self, i):
|
||||
if i < 0 or i >= self._len:
|
||||
raise IndexError('index out of range')
|
||||
|
||||
def __del__(self):
|
||||
if self.data_file:
|
||||
self.data_file.close()
|
||||
|
||||
# @lru_cache(maxsize=8)
|
||||
def __getitem__(self, idx):
|
||||
if not self.data_file:
|
||||
self.read_data(self.path)
|
||||
if isinstance(idx, int):
|
||||
i = idx
|
||||
self.check_index(i)
|
||||
tensor_size = self.sizes[self.dim_offsets[i]:self.dim_offsets[i + 1]]
|
||||
a = np.empty(tensor_size, dtype=self.dtype)
|
||||
self.data_file.seek(self.data_offsets[i] * self.element_size)
|
||||
self.data_file.readinto(a)
|
||||
return a
|
||||
elif isinstance(idx, slice):
|
||||
start, stop, step = idx.indices(len(self))
|
||||
if step != 1:
|
||||
raise ValueError("Slices into indexed_dataset must be contiguous")
|
||||
sizes = self.sizes[self.dim_offsets[start]:self.dim_offsets[stop]]
|
||||
size = sum(sizes)
|
||||
a = np.empty(size, dtype=self.dtype)
|
||||
self.data_file.seek(self.data_offsets[start] * self.element_size)
|
||||
self.data_file.readinto(a)
|
||||
offsets = list(accumulate(sizes))
|
||||
sents = np.split(a, offsets[:-1])
|
||||
return sents
|
||||
|
||||
def __len__(self):
|
||||
return self._len
|
||||
|
||||
def num_tokens(self, index):
|
||||
return self.sizes[index]
|
||||
|
||||
def size(self, index):
|
||||
return self.sizes[index]
|
||||
|
||||
@staticmethod
|
||||
def exists(path):
|
||||
return (os.path.exists(index_file_path(path)) and os.path.exists(data_file_path(path)))
|
||||
|
||||
@property
|
||||
def supports_prefetch(self):
|
||||
return False # avoid prefetching to save memory
|
||||
|
||||
|
||||
class IndexedCachedDataset(IndexedDataset):
|
||||
|
||||
def __init__(self, path):
|
||||
super().__init__(path)
|
||||
self.cache = None
|
||||
self.cache_index = {}
|
||||
|
||||
@property
|
||||
def supports_prefetch(self):
|
||||
return True
|
||||
|
||||
def prefetch(self, indices):
|
||||
if all(i in self.cache_index for i in indices):
|
||||
return
|
||||
if not self.data_file:
|
||||
self.read_data(self.path)
|
||||
indices = sorted(set(indices))
|
||||
total_size = 0
|
||||
for i in indices:
|
||||
total_size += self.data_offsets[i + 1] - self.data_offsets[i]
|
||||
self.cache = np.empty(total_size, dtype=self.dtype)
|
||||
ptx = 0
|
||||
self.cache_index.clear()
|
||||
for i in indices:
|
||||
self.cache_index[i] = ptx
|
||||
size = self.data_offsets[i + 1] - self.data_offsets[i]
|
||||
a = self.cache[ptx:ptx + size]
|
||||
self.data_file.seek(self.data_offsets[i] * self.element_size)
|
||||
self.data_file.readinto(a)
|
||||
ptx += size
|
||||
if self.data_file:
|
||||
# close and delete data file after prefetch so we can pickle
|
||||
self.data_file.close()
|
||||
self.data_file = None
|
||||
|
||||
# @lru_cache(maxsize=8)
|
||||
def __getitem__(self, idx):
|
||||
if isinstance(idx, int):
|
||||
i = idx
|
||||
self.check_index(i)
|
||||
tensor_size = self.sizes[self.dim_offsets[i]:self.dim_offsets[i + 1]]
|
||||
a = np.empty(tensor_size, dtype=self.dtype)
|
||||
ptx = self.cache_index[i]
|
||||
np.copyto(a, self.cache[ptx:ptx + a.size])
|
||||
return a
|
||||
elif isinstance(idx, slice):
|
||||
# Hack just to make this work, can optimizer later if necessary
|
||||
sents = []
|
||||
for i in range(*idx.indices(len(self))):
|
||||
sents.append(self[i])
|
||||
return sents
|
||||
|
||||
|
||||
class IndexedDatasetBuilder(object):
|
||||
|
||||
def __init__(self, out_file, dtype=np.int32):
|
||||
self.out_file = open(out_file, 'wb')
|
||||
self.dtype = dtype
|
||||
self.data_offsets = [0]
|
||||
self.dim_offsets = [0]
|
||||
self.sizes = []
|
||||
self.element_size = self.dtype().itemsize
|
||||
self.doc_idx = [0]
|
||||
|
||||
def add_item(self, tensor):
|
||||
bytes = self.out_file.write(np.array(tensor.numpy(), dtype=self.dtype))
|
||||
self.data_offsets.append(self.data_offsets[-1] + bytes / self.element_size)
|
||||
for s in tensor.size():
|
||||
self.sizes.append(s)
|
||||
self.dim_offsets.append(self.dim_offsets[-1] + len(tensor.size()))
|
||||
|
||||
def end_document(self):
|
||||
self.doc_idx.append(len(self.sizes))
|
||||
|
||||
def merge_file_(self, another_file):
|
||||
index = IndexedDataset(another_file)
|
||||
assert index.dtype == self.dtype
|
||||
|
||||
doc_offset = len(self.sizes)
|
||||
|
||||
begin = self.data_offsets[-1]
|
||||
for data_offset in index.data_offsets[1:]:
|
||||
self.data_offsets.append(begin + data_offset)
|
||||
self.sizes.extend(index.sizes)
|
||||
begin = self.dim_offsets[-1]
|
||||
for dim_offset in index.dim_offsets[1:]:
|
||||
self.dim_offsets.append(begin + dim_offset)
|
||||
self.doc_idx.extend((doc_offset + index.doc_idx)[1:])
|
||||
|
||||
with open(data_file_path(another_file), 'rb') as f:
|
||||
while True:
|
||||
data = f.read(1024)
|
||||
if data:
|
||||
self.out_file.write(data)
|
||||
else:
|
||||
break
|
||||
|
||||
def finalize(self, index_file):
|
||||
self.out_file.close()
|
||||
index = open(index_file, 'wb')
|
||||
index.write(b'TNTIDX\x00\x00')
|
||||
index.write(struct.pack('<Q', 1))
|
||||
index.write(struct.pack('<QQ', code(self.dtype), self.element_size))
|
||||
index.write(struct.pack('<QQ', len(self.data_offsets) - 1, len(self.sizes)))
|
||||
index.write(struct.pack('<Q', len(self.doc_idx)))
|
||||
write_longs(index, self.dim_offsets)
|
||||
write_longs(index, self.data_offsets)
|
||||
write_longs(index, self.sizes)
|
||||
write_longs(index, self.doc_idx)
|
||||
index.close()
|
||||
|
||||
|
||||
def _warmup_mmap_file(path):
|
||||
with open(path, 'rb') as stream:
|
||||
while stream.read(100 * 1024 * 1024):
|
||||
pass
|
||||
|
||||
|
||||
def exscan_from_cumsum_(arr):
|
||||
# given an array holding the result of an inclusive scan (cumsum),
|
||||
# convert to an exclusive scan (shift to the right)
|
||||
# [10, 30, 35, 50] --> [0, 10, 30, 35]
|
||||
if arr.size > 1:
|
||||
arr[1:] = arr[:-1]
|
||||
if arr.size > 0:
|
||||
arr[0] = 0
|
||||
|
||||
|
||||
def get_pointers_with_total(sizes, elemsize, dtype):
|
||||
"""Return a numpy array of type np.dtype giving the byte offsets.
|
||||
|
||||
Multiplies values in the sizes array by elemsize (bytes),
|
||||
and then computes an exclusive scan to get byte offsets.
|
||||
Returns the total number of bytes as second item in a tuple.
|
||||
"""
|
||||
|
||||
# scale values in sizes array by elemsize to get sizes in bytes
|
||||
pointers = np.array(sizes, dtype=dtype)
|
||||
pointers *= elemsize
|
||||
np.cumsum(pointers, axis=0, out=pointers)
|
||||
|
||||
# get total number of bytes from all sizes (last element)
|
||||
bytes_last = pointers[-1] if len(sizes) > 0 else 0
|
||||
|
||||
# convert to byte offsets
|
||||
exscan_from_cumsum_(pointers)
|
||||
|
||||
return pointers, bytes_last
|
||||
|
||||
|
||||
class MMapIndexedDataset(torch.utils.data.Dataset):
|
||||
|
||||
class Index(object):
|
||||
_HDR_MAGIC = b'MMIDIDX\x00\x00'
|
||||
|
||||
@classmethod
|
||||
def writer(cls, path, dtype):
|
||||
|
||||
class _Writer(object):
|
||||
|
||||
def __enter__(self):
|
||||
self._file = open(path, 'wb')
|
||||
|
||||
self._file.write(cls._HDR_MAGIC)
|
||||
self._file.write(struct.pack('<Q', 1))
|
||||
self._file.write(struct.pack('<B', code(dtype)))
|
||||
|
||||
return self
|
||||
|
||||
@staticmethod
|
||||
def _get_pointers(sizes, npdtype):
|
||||
"""Return a numpy array of byte offsets given a list of sizes.
|
||||
|
||||
Multiplies values in the sizes array by dtype size (bytes),
|
||||
and then computes an exclusive scan to get byte offsets.
|
||||
"""
|
||||
|
||||
# compute element sizes in bytes
|
||||
pointers, _ = get_pointers_with_total(sizes, dtype().itemsize, npdtype)
|
||||
return pointers
|
||||
|
||||
def write(self, sizes, doc_idx):
|
||||
self._file.write(struct.pack('<Q', len(sizes)))
|
||||
self._file.write(struct.pack('<Q', len(doc_idx)))
|
||||
|
||||
sizes32 = np.array(sizes, dtype=np.int32)
|
||||
self._file.write(sizes32.tobytes(order='C'))
|
||||
del sizes32
|
||||
|
||||
pointers = self._get_pointers(sizes, np.int64)
|
||||
del sizes
|
||||
self._file.write(pointers.tobytes(order='C'))
|
||||
del pointers
|
||||
|
||||
doc_idx = np.array(doc_idx, dtype=np.int64)
|
||||
self._file.write(doc_idx.tobytes(order='C'))
|
||||
|
||||
def __exit__(self, exc_type, exc_val, exc_tb):
|
||||
self._file.close()
|
||||
|
||||
return _Writer()
|
||||
|
||||
def __init__(self, path, skip_warmup=False):
|
||||
with open(path, 'rb') as stream:
|
||||
magic_test = stream.read(9)
|
||||
assert self._HDR_MAGIC == magic_test, ('Index file doesn\'t match expected format. '
|
||||
'Make sure that --dataset-impl is configured properly.')
|
||||
version = struct.unpack('<Q', stream.read(8))
|
||||
assert (1, ) == version
|
||||
|
||||
dtype_code, = struct.unpack('<B', stream.read(1))
|
||||
self._dtype = dtypes[dtype_code][0] #numpy type
|
||||
self._dtype_size = self._dtype().itemsize
|
||||
|
||||
self._len = struct.unpack('<Q', stream.read(8))[0]
|
||||
self._doc_count = struct.unpack('<Q', stream.read(8))[0]
|
||||
offset = stream.tell()
|
||||
|
||||
if not skip_warmup:
|
||||
print(" warming up index mmap file...")
|
||||
_warmup_mmap_file(path)
|
||||
|
||||
self._bin_buffer_mmap = np.memmap(path, mode='r', order='C')
|
||||
self._bin_buffer = memoryview(self._bin_buffer_mmap)
|
||||
print(" reading sizes...")
|
||||
self._sizes = np.frombuffer(self._bin_buffer, dtype=np.int32, count=self._len, offset=offset)
|
||||
print(" reading pointers...")
|
||||
self._pointers = np.frombuffer(self._bin_buffer,
|
||||
dtype=np.int64,
|
||||
count=self._len,
|
||||
offset=offset + self._sizes.nbytes)
|
||||
print(" reading document index...")
|
||||
self._doc_idx = np.frombuffer(self._bin_buffer,
|
||||
dtype=np.int64,
|
||||
count=self._doc_count,
|
||||
offset=offset + self._sizes.nbytes + self._pointers.nbytes)
|
||||
|
||||
def __del__(self):
|
||||
self._bin_buffer_mmap._mmap.close()
|
||||
del self._bin_buffer_mmap
|
||||
|
||||
@property
|
||||
def dtype(self):
|
||||
return self._dtype
|
||||
|
||||
@property
|
||||
def sizes(self):
|
||||
return self._sizes
|
||||
|
||||
@property
|
||||
def doc_idx(self):
|
||||
return self._doc_idx
|
||||
|
||||
@lru_cache(maxsize=8)
|
||||
def __getitem__(self, i):
|
||||
return self._pointers[i], self._sizes[i]
|
||||
|
||||
def __len__(self):
|
||||
return self._len
|
||||
|
||||
def __init__(self, path, skip_warmup=False):
|
||||
super().__init__()
|
||||
|
||||
self._path = None
|
||||
self._index = None
|
||||
self._bin_buffer = None
|
||||
|
||||
self._do_init(path, skip_warmup)
|
||||
|
||||
def __getstate__(self):
|
||||
return self._path
|
||||
|
||||
def __setstate__(self, state):
|
||||
self._do_init(state)
|
||||
|
||||
def _do_init(self, path, skip_warmup):
|
||||
self._path = path
|
||||
self._index = self.Index(index_file_path(self._path), skip_warmup)
|
||||
|
||||
if not skip_warmup:
|
||||
print(" warming up data mmap file...")
|
||||
_warmup_mmap_file(data_file_path(self._path))
|
||||
print(" creating numpy buffer of mmap...")
|
||||
self._bin_buffer_mmap = np.memmap(data_file_path(self._path), mode='r', order='C')
|
||||
print(" creating memory view of numpy buffer...")
|
||||
self._bin_buffer = memoryview(self._bin_buffer_mmap)
|
||||
|
||||
def __del__(self):
|
||||
self._bin_buffer_mmap._mmap.close()
|
||||
del self._bin_buffer_mmap
|
||||
del self._index
|
||||
|
||||
def __len__(self):
|
||||
return len(self._index)
|
||||
|
||||
# @lru_cache(maxsize=8)
|
||||
def __getitem__(self, idx):
|
||||
if isinstance(idx, int):
|
||||
ptr, size = self._index[idx]
|
||||
np_array = np.frombuffer(self._bin_buffer, dtype=self._index.dtype, count=size, offset=ptr)
|
||||
return np_array
|
||||
elif isinstance(idx, slice):
|
||||
start, stop, step = idx.indices(len(self))
|
||||
if step != 1:
|
||||
raise ValueError("Slices into indexed_dataset must be contiguous")
|
||||
ptr = self._index._pointers[start]
|
||||
sizes = self._index._sizes[idx]
|
||||
offsets = list(accumulate(sizes))
|
||||
total_size = sum(sizes)
|
||||
np_array = np.frombuffer(self._bin_buffer, dtype=self._index.dtype, count=total_size, offset=ptr)
|
||||
sents = np.split(np_array, offsets[:-1])
|
||||
return sents
|
||||
|
||||
def get(self, idx, offset=0, length=None):
|
||||
""" Retrieves a single item from the dataset with the option to only
|
||||
return a portion of the item.
|
||||
|
||||
get(idx) is the same as [idx] but get() does not support slicing.
|
||||
"""
|
||||
ptr, size = self._index[idx]
|
||||
if length is None:
|
||||
length = size - offset
|
||||
ptr += offset * np.dtype(self._index.dtype).itemsize
|
||||
np_array = np.frombuffer(self._bin_buffer, dtype=self._index.dtype, count=length, offset=ptr)
|
||||
return np_array
|
||||
|
||||
@property
|
||||
def sizes(self):
|
||||
return self._index.sizes
|
||||
|
||||
def size(self, index):
|
||||
return self._index.sizes[index]
|
||||
|
||||
@property
|
||||
def doc_idx(self):
|
||||
return self._index.doc_idx
|
||||
|
||||
def get_doc_idx(self):
|
||||
return self._index._doc_idx
|
||||
|
||||
def set_doc_idx(self, doc_idx_):
|
||||
self._index._doc_idx = doc_idx_
|
||||
|
||||
@property
|
||||
def supports_prefetch(self):
|
||||
return False
|
||||
|
||||
@staticmethod
|
||||
def exists(path):
|
||||
return (os.path.exists(index_file_path(path)) and os.path.exists(data_file_path(path)))
|
||||
|
||||
@property
|
||||
def dtype(self):
|
||||
return self._index.dtype
|
||||
|
||||
|
||||
class MMapIndexedDatasetBuilder(object):
|
||||
|
||||
def __init__(self, out_file, dtype=np.int64):
|
||||
self._data_file = open(out_file, 'wb')
|
||||
self._dtype = [np_dt for np_dt, torch_dt in dtypes.values() if dtype in [np_dt, torch_dt]][0]
|
||||
self._sizes = []
|
||||
self._doc_idx = [0]
|
||||
|
||||
def add_item(self, tensor):
|
||||
""" write the tensor to the file and update its size in the index"""
|
||||
np_array = np.array(tensor.numpy(), dtype=self._dtype)
|
||||
self._data_file.write(np_array.tobytes(order='C'))
|
||||
self._sizes.append(np_array.size)
|
||||
|
||||
def add_items(self, arr_list):
|
||||
""" write a list of arrays to the file and update their sizes in the index"""
|
||||
np_arrays = [arr.astype(self._dtype) for arr in arr_list]
|
||||
self._data_file.writelines([arr.tobytes(order='C') for arr in np_arrays])
|
||||
for arr in np_arrays:
|
||||
self._sizes.append(arr.size)
|
||||
|
||||
def add_item_numpy(self, np_array):
|
||||
if np_array.dtype != self._dtype:
|
||||
np_array = np_array.astype(self._dtype)
|
||||
self._data_file.write(np_array.tobytes(order='C'))
|
||||
self._sizes.append(np_array.size)
|
||||
|
||||
def end_document(self):
|
||||
self._doc_idx.append(len(self._sizes))
|
||||
|
||||
def merge_file_(self, another_file):
|
||||
# Concatenate index
|
||||
index = MMapIndexedDataset.Index(index_file_path(another_file))
|
||||
assert index.dtype == self._dtype
|
||||
|
||||
total_len = len(index.sizes) + len(self._sizes)
|
||||
print(f" concat {another_file} size={len(index.sizes)} for a total size of {total_len}")
|
||||
|
||||
offset = len(self._sizes)
|
||||
self._sizes.extend(index.sizes)
|
||||
self._doc_idx.extend((offset + index.doc_idx)[1:])
|
||||
|
||||
# Concatenate data
|
||||
with open(data_file_path(another_file), 'rb') as f:
|
||||
shutil.copyfileobj(f, self._data_file)
|
||||
self._data_file.flush()
|
||||
assert os.stat(self._data_file.name).st_size != 0, f"Zero-sized file: {self._data_file.name}"
|
||||
|
||||
def finalize(self, index_file):
|
||||
self._data_file.close()
|
||||
|
||||
with MMapIndexedDataset.Index.writer(index_file, self._dtype) as index:
|
||||
index.write(self._sizes, self._doc_idx)
|
||||
@@ -0,0 +1,52 @@
|
||||
# Copyright (c) Microsoft Corporation.
|
||||
# SPDX-License-Identifier: Apache-2.0
|
||||
|
||||
# DeepSpeed Team
|
||||
|
||||
import numpy as np
|
||||
|
||||
from deepspeed.utils import logger
|
||||
from .indexed_dataset import MMapIndexedDatasetBuilder
|
||||
|
||||
|
||||
def find_fit_int_dtype(min_value, max_value):
|
||||
if min_value >= 0:
|
||||
if max_value <= 255:
|
||||
return np.uint8
|
||||
elif max_value <= 65535:
|
||||
return np.uint16
|
||||
elif max_value <= 4294967295:
|
||||
return np.uint32
|
||||
else:
|
||||
return np.uint64
|
||||
else:
|
||||
if max_value <= 127 and min_value >= -128:
|
||||
return np.int8
|
||||
elif max_value <= 32767 and min_value >= -32768:
|
||||
return np.int16
|
||||
elif max_value <= 2147483647 and min_value >= -2147483648:
|
||||
return np.int32
|
||||
else:
|
||||
return np.int64
|
||||
|
||||
|
||||
def split_index(start_idx, end_idx, num_partitions):
|
||||
partition_boundaries = np.linspace(start_idx, end_idx, dtype=int, num=num_partitions + 1)
|
||||
return [(partition_boundaries[i], partition_boundaries[i + 1]) for i in range(num_partitions)]
|
||||
|
||||
|
||||
def split_dataset(dataset, num_workers, worker_id, num_threads):
|
||||
worker_splits = split_index(0, len(dataset), num_workers)
|
||||
thread_splits = split_index(worker_splits[worker_id][0], worker_splits[worker_id][1], num_threads)
|
||||
return worker_splits, thread_splits
|
||||
|
||||
|
||||
def create_mmap_dataset_builder(fname, dtype):
|
||||
logger.info(f"Creating mmap dataset builder at {fname}.")
|
||||
return MMapIndexedDatasetBuilder(f"{fname}.bin", dtype=dtype)
|
||||
|
||||
|
||||
def close_mmap_dataset_builder(builder, fname):
|
||||
builder.end_document()
|
||||
builder.finalize(f"{fname}.idx")
|
||||
logger.info(f"Finalized mmap dataset builder at {fname}.")
|
||||
@@ -0,0 +1,492 @@
|
||||
# Copyright (c) Microsoft Corporation.
|
||||
# SPDX-License-Identifier: Apache-2.0
|
||||
|
||||
# DeepSpeed Team
|
||||
|
||||
# support/questions/maintenance: github user @brunomaga or @deepspeedai/deepspeed
|
||||
|
||||
import random
|
||||
import os
|
||||
import math
|
||||
import numpy as np
|
||||
from torch.optim.lr_scheduler import LRScheduler
|
||||
from torch.optim.optimizer import Optimizer
|
||||
from torch.utils.data import DataLoader, DistributedSampler
|
||||
from deepspeed.utils import logger
|
||||
from deepspeed.runtime.pipe.engine import PipelineEngine
|
||||
from deepspeed.runtime.data_pipeline.constants import *
|
||||
from deepspeed.runtime.data_pipeline.data_sampling.indexed_dataset import MMapIndexedDataset
|
||||
from deepspeed.runtime.data_pipeline.data_sampling.data_analyzer import DistributedDataAnalyzer
|
||||
import pathlib
|
||||
|
||||
|
||||
def batch_by_seqlens(
|
||||
seqlens,
|
||||
max_tokens,
|
||||
sequence_ids_per_mb=None,
|
||||
min_batch_size=1,
|
||||
max_batch_size=None,
|
||||
sequence_picking_order="dataloader",
|
||||
effective_batch_size=1,
|
||||
required_microbatches_of_same_size=False,
|
||||
verbose=False,
|
||||
seed=None,
|
||||
):
|
||||
"""
|
||||
Yield mini-batches of indices bucketed by size. Batches may contain sequences of different lengths.
|
||||
Similar to "Attention is all you need", Section 5.1:
|
||||
"sequence pairs were batched together by approximate sequence length. Each training batch
|
||||
contained a set of sequence pairs containing approximately X source tokens and X target tokens"
|
||||
|
||||
Arguments:
|
||||
- `seqlens`: a list of difficulties (metric values) for every sample in the dataset;
|
||||
- `max_tokens`: maximum cap in total difficulty in a batch;
|
||||
- `min_batch_size`: smallest allowed size of a batch;
|
||||
- `min_batch_size`: largest allowed size of a batch;
|
||||
- `sequence_picking_order`: order in which to process samples: "dataloader" (default), "random" or "seqlen" (ascending)
|
||||
- `effective_batch_size`: effective batch size;
|
||||
- `required_microbatches_of_same_size`: enable if each mini-batch (in a total of `batch_size_multiple`
|
||||
micro-batches per batch), should have all micro-batches with the same batch size ie the same
|
||||
number of sequences.
|
||||
- `verbose`: print debug information;
|
||||
- `seed`: random seed for reproducibility;
|
||||
|
||||
Returns:
|
||||
- `microbatch_ids`: list of tuple of batch id and samples ids per microbatch
|
||||
- `batch_sizes`: the effective batch size of each batch, used for to compute the scaled LR
|
||||
- `batch_max_seqlens`: the max seqlen across all microbatches in a batch
|
||||
"""
|
||||
|
||||
assert sequence_picking_order in ["random", "seqlen", "dataloader"]
|
||||
if sequence_ids_per_mb is None:
|
||||
metrics = list(zip(seqlens, range(len(seqlens)))) # use all samples
|
||||
else:
|
||||
metrics = list(zip(np.array(seqlens)[sequence_ids_per_mb], sequence_ids_per_mb))
|
||||
|
||||
if sequence_picking_order == 'random':
|
||||
metric_random = random.Random(seed)
|
||||
metric_random.shuffle(metrics)
|
||||
if sequence_picking_order == 'seqlen':
|
||||
metrics = sorted(metrics)
|
||||
|
||||
# go through metrics, warn user, and filter samples that alone exceed the max batch threshold
|
||||
long_ids = [idx for val, idx in metrics if val > max_tokens]
|
||||
if len(long_ids) > 0:
|
||||
logger.warning(f"Data indices {long_ids} ignored as metrics exceed {max_tokens}.")
|
||||
logger.info(f"Original dataset length: {len(metrics)}. New dataset length: {len(long_ids)}")
|
||||
metrics = [m for m in metrics if m[1] not in long_ids]
|
||||
|
||||
def is_microbatch_valid(metrics):
|
||||
if min_batch_size and len(metrics) < min_batch_size: return False # insufficient sample count
|
||||
if max_batch_size and len(metrics) > max_batch_size: return False # too many samples
|
||||
if sum([m[0] for m in metrics]) > max_tokens: return False # exceeds max
|
||||
return True
|
||||
|
||||
# go through all samples and pack then in microbatches of metric sums below the threshold
|
||||
# `required_microbatches_of_same_size` means all minibatches in a batch must be of equal size
|
||||
equal_size_multiple = effective_batch_size if required_microbatches_of_same_size else 1
|
||||
microbatches = []
|
||||
batch_init = 0
|
||||
while batch_init < len(metrics):
|
||||
|
||||
# we iterate over possible effective batch sizes (groups of microbatches of same size)
|
||||
valid_batch_end = batch_init
|
||||
for batch_end in range(batch_init + equal_size_multiple, len(metrics), equal_size_multiple):
|
||||
|
||||
# attempt effective batch
|
||||
batch = metrics[batch_init:batch_end]
|
||||
|
||||
# pick interleaved samples for each microbatch to help with load balancing
|
||||
# (in the ordered use case), and to replicate what the distributed sampler does.
|
||||
mbs = [batch[b::equal_size_multiple] for b in range(equal_size_multiple)]
|
||||
|
||||
# if they are all valid micro-batches, keep them until you find longer mbatches, if any
|
||||
is_batch_valid = all([is_microbatch_valid(mb) for mb in mbs])
|
||||
if is_batch_valid:
|
||||
valid_batch_end = batch_end
|
||||
|
||||
if batch_init == valid_batch_end: break # last batch is not valid (size zero), so we are done
|
||||
batch = metrics[batch_init:valid_batch_end]
|
||||
mbs = [batch[b::equal_size_multiple] for b in range(equal_size_multiple)]
|
||||
batch_init += sum([len(l) for l in mbs])
|
||||
microbatches += mbs
|
||||
|
||||
# make sure we give the same number of (micro-)batches to each dataloader by trimming the dataset
|
||||
assert len(microbatches) >= effective_batch_size, "not enough datapoints to create a single sample per dataloader"
|
||||
microbatches = microbatches[:len(microbatches) - len(microbatches) % effective_batch_size]
|
||||
|
||||
#compute the effective batch size for each microbatch.
|
||||
batch_sizes, batch_max_seqlens, microbatch_ids = [], [], []
|
||||
for rank in range(0, len(microbatches), effective_batch_size):
|
||||
batch_id = rank // effective_batch_size
|
||||
mbs = microbatches[rank:rank + effective_batch_size]
|
||||
# compute the number of samples (not tokens) in this batch (not microbatch)
|
||||
n_sequences = sum([len(mb) for mb in mbs])
|
||||
# compute the longest sequence (as number of tokens) in this batch (not microbatch)
|
||||
sequence_ids_per_mb = [[m[1] for m in metrics] for metrics in mbs]
|
||||
sequence_lens_per_mb = [[m[0] for m in metrics] for metrics in mbs]
|
||||
batch_max_seqlen = max([max(seqlens) for seqlens in sequence_lens_per_mb])
|
||||
batch_and_mb_ids = zip([batch_id] * effective_batch_size, sequence_ids_per_mb)
|
||||
batch_sizes.append(n_sequences)
|
||||
batch_max_seqlens.append(batch_max_seqlen)
|
||||
microbatch_ids += batch_and_mb_ids
|
||||
if verbose:
|
||||
n_tokens_per_mb = [sum([m[0] for m in mb]) for mb in mbs]
|
||||
n_sequences_per_mb = [len(mb) for mb in mbs]
|
||||
assert all([n <= max_tokens for n in n_tokens_per_mb]), "size of microbatch exceeds max tokens"
|
||||
logger.info(
|
||||
f"Batch id {batch_id} contains in total {len(mbs)} microbatches or {n_sequences} sequences. "\
|
||||
f"n_sequences per microbatch {n_sequences_per_mb}. "\
|
||||
f"n_tokens per microbatch {n_tokens_per_mb}. "\
|
||||
f"sequence ids per microbatch: {sequence_ids_per_mb}. "\
|
||||
f"sequence lengths per microbatch: {sequence_lens_per_mb}.")
|
||||
|
||||
# return the sample ids of each microbatch, and the batch sizes
|
||||
assert len(batch_sizes) == len(microbatch_ids) // effective_batch_size
|
||||
return microbatch_ids, batch_sizes, batch_max_seqlens
|
||||
|
||||
|
||||
def scale_lr(base_batch_size, batch_size, base_lr=1, method="linear"):
|
||||
""" given a reference lr and batch_size, compute the new LR for a given batch size """
|
||||
if method == "linear":
|
||||
# Linear Scaling Rule: "When the minibatch size is multiplied by k, multiply the learning
|
||||
# rate by k" (Accurate, Large Minibatch SGD: Training ImageNet in 1 Hour, Goyal et al)
|
||||
return base_lr * batch_size / base_batch_size
|
||||
if method == "sqrt":
|
||||
# Square Root scaling: "when multiplying the batch size by k, multiply the learning rate
|
||||
# by √k, to keep the variance in the gradient expectation constant"
|
||||
# (A. Krizhevsky. One weird trick for parallelizing convolutional neural networks)
|
||||
return base_lr * math.sqrt(batch_size / base_batch_size)
|
||||
elif method == None or method.upper() == "NONE":
|
||||
return base_lr
|
||||
raise ValueError("Unknown scaling method: {}".format(method))
|
||||
|
||||
|
||||
def dataloader_for_variable_batch_size(
|
||||
dataset,
|
||||
microbatch_ids,
|
||||
batch_max_seqlens,
|
||||
dataloader_rank=0,
|
||||
dataloader_batch_size=1,
|
||||
dataloader_num_replicas=1,
|
||||
dataloader_collate_fn=None,
|
||||
dataloader_num_workers=2,
|
||||
dataloader_pin_memory=False,
|
||||
required_microbatches_of_same_seqlen=False,
|
||||
sample_padding_fn=None,
|
||||
):
|
||||
|
||||
# equidistantly distribute the microbatches across the replicas in an interleaved fashion.
|
||||
sampler = DistributedSampler(
|
||||
dataset=microbatch_ids,
|
||||
num_replicas=dataloader_num_replicas,
|
||||
rank=dataloader_rank,
|
||||
shuffle=False,
|
||||
drop_last=False,
|
||||
)
|
||||
|
||||
# collate function wraps user-defined collate function to the variable batch data
|
||||
def collate_fn_wrapper(list_microbatch_ids):
|
||||
# each batch is a list of sample ids that fill up to the max tokens per batch
|
||||
# we return the collated batch of all dataset samples of all input batches.
|
||||
batch = []
|
||||
for batch_id, microbatch_ids in list_microbatch_ids:
|
||||
batch_data = [dataset[idx] for idx in microbatch_ids]
|
||||
if required_microbatches_of_same_seqlen:
|
||||
assert sample_padding_fn is not None, \
|
||||
"padding dataloader_padding_fn must be provided if required_microbatches_of_same_seqlen is True"
|
||||
max_seqlen = batch_max_seqlens[batch_id]
|
||||
assert all([len(sample) <= max_seqlen for sample in batch_data]), \
|
||||
"some samples are longer than the computed max seqlen for the batch those samples belong to"
|
||||
batch_data = [sample_padding_fn(sample, max_seqlen) for sample in batch_data]
|
||||
batch += batch_data
|
||||
return dataloader_collate_fn(batch) if dataloader_collate_fn else batch
|
||||
|
||||
dataloader = DataLoader(
|
||||
dataset=microbatch_ids,
|
||||
batch_size=dataloader_batch_size,
|
||||
sampler=sampler,
|
||||
num_workers=dataloader_num_workers,
|
||||
collate_fn=collate_fn_wrapper,
|
||||
pin_memory=dataloader_pin_memory,
|
||||
)
|
||||
|
||||
deepspeed_io_kwargs = dict(
|
||||
dataset=microbatch_ids,
|
||||
batch_size=dataloader_batch_size,
|
||||
pin_memory=dataloader_pin_memory,
|
||||
data_sampler=sampler,
|
||||
collate_fn=collate_fn_wrapper,
|
||||
num_local_io_workers=dataloader_num_workers,
|
||||
)
|
||||
|
||||
return dataloader, deepspeed_io_kwargs
|
||||
|
||||
|
||||
class VariableBatchSizeLR(LRScheduler):
|
||||
""" an LR scheduler that scales the LR of a given scheduler's LR """
|
||||
|
||||
@property
|
||||
def optimizer(self):
|
||||
return self.base_lr_scheduler.optimizer
|
||||
|
||||
def __init__(self,
|
||||
lr_scheduler,
|
||||
base_batch_size,
|
||||
batch_sizes,
|
||||
dataloader,
|
||||
lr_scaling_method="linear",
|
||||
last_epoch=-1,
|
||||
verbose=False):
|
||||
self.batch_sizes = batch_sizes
|
||||
self.base_batch_size = base_batch_size
|
||||
self.lr_scaling_method = lr_scaling_method
|
||||
self.dataloader = dataloader
|
||||
self.base_lr_scheduler = lr_scheduler
|
||||
# the following exist in LRScheduler but not in DeepSpeed's LRScheduler so we redefine them here
|
||||
self.base_lrs = self.base_lr_scheduler.get_lr()
|
||||
self.last_epoch = last_epoch
|
||||
self.verbose = verbose
|
||||
self.step(0) # scale LR for first sample in the dataloader
|
||||
|
||||
def state_dict(self):
|
||||
return {
|
||||
'base_lr_scheduler': self.base_lr_scheduler.state_dict()
|
||||
} | {
|
||||
'base_batch_size': self.base_batch_size,
|
||||
'lr_scaling_method': self.lr_scaling_method,
|
||||
'batch_sizes': self.batch_sizes,
|
||||
'base_lrs': self.base_lrs,
|
||||
'last_epoch': self.last_epoch,
|
||||
'verbose': self.verbose,
|
||||
}
|
||||
|
||||
def load_state_dict(self, state_dict):
|
||||
self.base_lr_scheduler.load_state_dict(state_dict['base_lr_scheduler'])
|
||||
self.base_batch_size = state_dict['base_batch_size']
|
||||
self.lr_scaling_method = state_dict['lr_scaling_method']
|
||||
self.batch_sizes = state_dict['batch_sizes']
|
||||
self.base_lrs = state_dict['base_lrs']
|
||||
self.last_epoch = state_dict['last_epoch']
|
||||
self.verbose = state_dict['verbose']
|
||||
|
||||
def get_last_lr(self):
|
||||
return self.base_lr_scheduler._last_lr
|
||||
|
||||
def get_lr(self):
|
||||
return [group['lr'] for group in self.base_lr_scheduler.optimizer.param_groups]
|
||||
|
||||
def step(self, epoch=None):
|
||||
# call the base scheduler's step method to get LR for next epoch
|
||||
# Note: optimizer.step precedes lr_scheduler.step(), so the stepping workflow is:
|
||||
# init: lr_scheduler.step(0) --> set LR for epoch 0
|
||||
# epoch 0: optimizer.step(); lr_scheduler.step(1) --> set LR for epoch 1
|
||||
# epoch 1: optimizer.step(); lr_scheduler.step(2) --> set LR for epoch 2
|
||||
|
||||
# reset unscaled LRs (to the original scheduler's one) to be able to step the base LR scheduler
|
||||
# Note: epoch==0: reset LR scheduler; epoch==None: scale LR for next epoch;
|
||||
unscaled_lrs = self.base_lrs if epoch == 0 else self.get_last_lr()
|
||||
for group, lr in zip(self.base_lr_scheduler.optimizer.param_groups, unscaled_lrs):
|
||||
group['lr'] = lr
|
||||
|
||||
self.base_lr_scheduler.step(epoch) # set unscaled lr, _step_count, last_epoch, _last_lr for new epoch
|
||||
|
||||
# scale the learning rate for the next iteration for each parameter group.
|
||||
self.last_epoch = self.last_epoch + 1 if epoch is None else epoch
|
||||
# batch sizes are precomputed and stored in batch_sizes se we loop around to get the next one
|
||||
batch_size = self.batch_sizes[self.last_epoch % len(self.batch_sizes)]
|
||||
for group in self.base_lr_scheduler.optimizer.param_groups:
|
||||
group['lr'] = scale_lr(self.base_batch_size, batch_size, group['lr'], self.lr_scaling_method)
|
||||
|
||||
if self.verbose:
|
||||
logger.info(
|
||||
f"Next batch id {self.last_epoch}. "\
|
||||
f"Reference batch_size {self.base_batch_size} and lr {unscaled_lrs}. "\
|
||||
f"Scaled batch_size {batch_size} and lr {self.get_lr()}.")
|
||||
|
||||
|
||||
def lr_scheduler_for_variable_batch_size(base_batch_size,
|
||||
batch_sizes,
|
||||
dataloader,
|
||||
lr_scheduler_or_optimizer,
|
||||
lr_scaling_method='linear',
|
||||
verbose=False):
|
||||
"""
|
||||
returns a class that provides an LR scheduler that scales the learning rate at every
|
||||
iteration taking into account the batch size of that iteration.
|
||||
If learning rate is constant, ie no LR scheduler, then the base LR will be taken from the
|
||||
constant LR values in the optimizer param groups. Otherwise from the scheduler's LR.
|
||||
|
||||
Arguments:
|
||||
- `base_batch_size`: the batch size that the base LR in the optimizer or scheduler refers to;
|
||||
- `lr_scaling_method`: method to use to scale LR - see `scale_lr()`;
|
||||
- `lr_scheduler_or_optimizer`: one instance of `LRScheduler` or `Optimizer` to be used as base;
|
||||
- `batch_sizes`: the effective batch size of each batch in the dataloader;
|
||||
|
||||
Returns the new LRScheduler
|
||||
"""
|
||||
|
||||
class StubLRScheduler(LRScheduler):
|
||||
""" a stub LR scheduler that does not change the LR, keeps it constant """
|
||||
|
||||
def get_lr(self) -> float:
|
||||
return self.base_lrs
|
||||
|
||||
if isinstance(lr_scheduler_or_optimizer, Optimizer):
|
||||
lr_scheduler = StubLRScheduler(lr_scheduler_or_optimizer)
|
||||
elif hasattr(lr_scheduler_or_optimizer, 'optimizer'): #LRScheduler or DeepSpeed 'object' schedulers
|
||||
assert isinstance(lr_scheduler_or_optimizer.optimizer, Optimizer)
|
||||
lr_scheduler = lr_scheduler_or_optimizer
|
||||
else:
|
||||
raise ValueError("Unknown type for lr_scheduler_or_optimizer: {}".format(type(lr_scheduler_or_optimizer)))
|
||||
|
||||
return VariableBatchSizeLR(lr_scheduler=lr_scheduler,
|
||||
base_batch_size=base_batch_size,
|
||||
batch_sizes=batch_sizes,
|
||||
dataloader=dataloader,
|
||||
lr_scaling_method=lr_scaling_method,
|
||||
verbose=verbose)
|
||||
|
||||
|
||||
def get_dataloader_and_lr_scheduler_for_variable_batch_size_deepspeed(dataset,
|
||||
engine,
|
||||
dataset_seqlens=None,
|
||||
dataset_filter_ids=None,
|
||||
dataloader_collate_fn=None,
|
||||
sample_padding_fn=None,
|
||||
batch_seqlens_fn=None):
|
||||
"""
|
||||
a simplified call to get_dataloader_and_lr_scheduler_for_variable_batch_size for the deepspeed runtime.
|
||||
Needs the seqlens of every sample. It will try three alternatives:
|
||||
- if `dataset_seqlens` is provided by user, use that.
|
||||
- otherwise, looks for the seqlen metric path (in the connfig) that contains the output of the Data Analyzer
|
||||
- otherwise, use the user-provided function `batch_seqlens_fn` and call Data Analyzer to output seqlen metric
|
||||
See `batch_by_seqlens()` for arguments and more documentation.
|
||||
"""
|
||||
data_efficiency_config = engine._config.data_efficiency_config
|
||||
data_sampling_config = data_efficiency_config[DATA_SAMPLING]
|
||||
batching_config = data_sampling_config[DYNAMIC_BATCHING]
|
||||
assert batching_config[DYNAMIC_BATCHING_ENABLED], "Dynamic batching is not enabled in the config"
|
||||
|
||||
if dataset_seqlens is None:
|
||||
# In seqlen provided by user, look for the seqlen metric that was output by the Data Analyzer
|
||||
# (see the main in deepspeed/runtime/data_pipeline/data_sampling/data_analyzer.py for an example)
|
||||
metrics_path = batching_config[DYNAMIC_BATCHING_METRICS_PATH]
|
||||
sample_to_seqlen_path = os.path.join(metrics_path, "seqlen/seqlen_sample_to_metric")
|
||||
if not (os.path.exists(f"{sample_to_seqlen_path}.bin") and os.path.exists(f"{sample_to_seqlen_path}.idx")):
|
||||
# if the metric files are not found, we run the DataAnalyzer to write the metric files
|
||||
msg = f"Cannot find metric files for sequence length in {sample_to_seqlen_path}.idx or *.bin."
|
||||
msg += " We will run data analyzer to generated them..."
|
||||
logger.warning(msg)
|
||||
|
||||
if batch_seqlens_fn is None:
|
||||
raise ValueError("sample_seqlen_fn must be provided if dataset_seqlens is not provided")
|
||||
|
||||
DistributedDataAnalyzer(
|
||||
dataset=dataset,
|
||||
metric_functions=[batch_seqlens_fn],
|
||||
collate_fn=dataloader_collate_fn,
|
||||
batch_size=2**10, # batch size for map-reduce, not training
|
||||
num_workers=engine.world_size,
|
||||
worker_id=engine.global_rank,
|
||||
save_path=pathlib.Path(metrics_path),
|
||||
metric_types=['single_value_per_sample'],
|
||||
metric_names=["seqlen"],
|
||||
device=engine.device,
|
||||
).run_map_reduce()
|
||||
|
||||
dataset_seqlens = MMapIndexedDataset(sample_to_seqlen_path, skip_warmup=True)
|
||||
assert len(dataset_seqlens) == len(dataset), \
|
||||
"Seqlens size does not match the input dataset size. If you changed the dataset, delete the metrics_path folder."
|
||||
|
||||
# TODO we are copying all seqlens into memory, we should adapt the code to use an iterative streamer
|
||||
# and use the other files output by DataAnalyzer that returns an ordered dictionary of seqlen to sample ids
|
||||
dataset_seqlens = np.array(list(dataset_seqlens), dtype=np.int64).flatten() # from Nx1 to N
|
||||
|
||||
dataloader, lr_scheduler, deepspeed_io_kwargs = get_dataloader_and_lr_scheduler_for_variable_batch_size(
|
||||
dataset=dataset,
|
||||
dataset_filter_ids=dataset_filter_ids,
|
||||
dataset_seqlens=dataset_seqlens,
|
||||
effective_batch_size=engine.train_batch_size(),
|
||||
max_tokens=batching_config[DYNAMIC_BATCHING_MAX_TOKENS],
|
||||
lr_scaling_method=batching_config[DYNAMIC_BATCHING_LR_SCALING_METHOD],
|
||||
sequence_picking_order=batching_config[DYNAMIC_BATCHING_SEQUENCE_PICKING_ORDER],
|
||||
min_batch_size=batching_config[DYNAMIC_BATCHING_MIN_BATCH_SIZE],
|
||||
max_batch_size=batching_config[DYNAMIC_BATCHING_MAX_BATCH_SIZE],
|
||||
dataloader_batch_size=engine.train_micro_batch_size_per_gpu(),
|
||||
dataloader_rank=engine.data_parallel_group.rank(),
|
||||
dataloader_num_replicas=engine.data_parallel_group.size(),
|
||||
dataloader_num_workers=data_sampling_config[DATA_SAMPLING_NUM_WORKERS],
|
||||
dataloader_collate_fn=dataloader_collate_fn,
|
||||
dataloader_pin_memory=data_sampling_config[DATA_SAMPLING_PIN_MEMORY],
|
||||
sample_padding_fn=sample_padding_fn,
|
||||
lr_scheduler_or_optimizer=engine.lr_scheduler or engine.optimizer,
|
||||
required_microbatches_of_same_size=isinstance(engine, PipelineEngine),
|
||||
required_microbatches_of_same_seqlen=isinstance(engine, PipelineEngine),
|
||||
verbose=batching_config[DYNAMIC_BATCHING_VERBOSE],
|
||||
seed=data_efficiency_config[DATA_EFFICIENCY_SEED],
|
||||
)
|
||||
return dataloader, lr_scheduler, deepspeed_io_kwargs
|
||||
|
||||
|
||||
def get_dataloader_and_lr_scheduler_for_variable_batch_size(
|
||||
dataset,
|
||||
dataset_seqlens,
|
||||
max_tokens,
|
||||
effective_batch_size,
|
||||
dataset_filter_ids=None,
|
||||
lr_scaling_method="linear",
|
||||
min_batch_size=1,
|
||||
max_batch_size=None,
|
||||
sequence_picking_order="dataloader",
|
||||
dataloader_batch_size=1,
|
||||
dataloader_rank=0,
|
||||
dataloader_num_replicas=1,
|
||||
dataloader_num_workers=0,
|
||||
dataloader_collate_fn=None,
|
||||
dataloader_pin_memory=False,
|
||||
lr_scheduler_or_optimizer=None,
|
||||
required_microbatches_of_same_size=False,
|
||||
required_microbatches_of_same_seqlen=False,
|
||||
sample_padding_fn=None,
|
||||
verbose=False,
|
||||
seed=None,
|
||||
):
|
||||
""" returns a dataloader and LR scheduler for the variable batch size. see `batch_by_seqlens()` for details. """
|
||||
|
||||
# effective_batch_size = train_micro_batch_size_per_gpu * gradient_accumulation_steps * number of dataloaders
|
||||
microbatch_ids, batch_sizes, batch_max_seqlens = batch_by_seqlens(
|
||||
seqlens=dataset_seqlens,
|
||||
max_tokens=max_tokens,
|
||||
sequence_ids_per_mb=dataset_filter_ids,
|
||||
min_batch_size=min_batch_size,
|
||||
max_batch_size=max_batch_size,
|
||||
sequence_picking_order=sequence_picking_order,
|
||||
effective_batch_size=effective_batch_size,
|
||||
required_microbatches_of_same_size=required_microbatches_of_same_size,
|
||||
verbose=verbose,
|
||||
seed=seed,
|
||||
)
|
||||
|
||||
dataloader, deepspeed_io_kwargs = dataloader_for_variable_batch_size(
|
||||
dataset=dataset,
|
||||
microbatch_ids=microbatch_ids,
|
||||
batch_max_seqlens=batch_max_seqlens,
|
||||
dataloader_rank=dataloader_rank,
|
||||
dataloader_num_replicas=dataloader_num_replicas,
|
||||
dataloader_batch_size=dataloader_batch_size,
|
||||
dataloader_collate_fn=dataloader_collate_fn,
|
||||
dataloader_num_workers=dataloader_num_workers,
|
||||
dataloader_pin_memory=dataloader_pin_memory,
|
||||
required_microbatches_of_same_seqlen=required_microbatches_of_same_seqlen,
|
||||
sample_padding_fn=sample_padding_fn,
|
||||
)
|
||||
|
||||
lr_scheduler = lr_scheduler_for_variable_batch_size(base_batch_size=effective_batch_size,
|
||||
batch_sizes=batch_sizes,
|
||||
lr_scaling_method=lr_scaling_method,
|
||||
lr_scheduler_or_optimizer=lr_scheduler_or_optimizer,
|
||||
dataloader=dataloader,
|
||||
verbose=verbose)
|
||||
|
||||
return dataloader, lr_scheduler, deepspeed_io_kwargs
|
||||
Reference in New Issue
Block a user