seed: int

dataset_path: str

batch_max_length: int

drop_last_batch: bool

target_only: bool

rank: int

tokenizer_info: dict[str, Any] = {}

vocab_size: Optional[int] = None

max_seq_len: int

inst_start: np.ndarray

inst_len: np.ndarray

resp_start: np.ndarray

def __init__(self, config: V1DatasetConfig):

super().__init__()

self.config = config

self.metadata = self._load_metadata()

# State

self._data: Optional[np.ndarray] = None

self._data_indices: Optional[V1DatasetIndices] = None

self._sampler: Optional[MultipackDistributedBatchSampler] = None

self._epoch = 0

def _load_metadata(self) -> V1DatasetMeta:

with open(os.path.join(self.config.dataset_path, "metadata.json"), "r") as f:

metadata = V1DatasetMeta(**json.load(f))

# Account for autoregressive shift

metadata.max_seq_len -= 1

# Compute vocab size from tokenizer info

assert metadata.vocab_size is None

metadata.vocab_size = find_multiple(metadata.tokenizer_info.pop("vocab_size"), 256)

return metadata

def _load_dataset_before_epoch_begin(self):

# Load tokens (only if not loaded)

if self._data is None:

self._data = np.load(os.path.join(self.config.dataset_path, "tokens.npy"), mmap_mode="r")

# Load indices

self._data_indices = V1DatasetIndices(**{f.name: np.load(os.path.join(self.config.dataset_path, f"epoch_{self._epoch}", f"{f.name}.npy"), mmap_mode="r")

for f in fields(V1DatasetIndices)})

self._epoch += 1

# Re-create sampler

self._sampler = MultipackDistributedBatchSampler(

lengths=self._data_indices.inst_len + self._data_indices.resp_len - 1, # Account for autoregressive shift

batch_max_length=self.config.batch_max_length,

drop_last_batch=self.config.drop_last_batch,

rank=self.config.rank,

num_replicas=self.config.num_replicas

)

def _load_batch(self, indices: np.ndarray):

# Load instructions and responses

assert self._data is not None and self._data_indices is not None

raw = {k: [] for k in ["inst", "resp"]}

for k, v in raw.items():

start = getattr(self._data_indices, f"{k}_start")[indices]

end = start + getattr(self._data_indices, f"{k}_len")[indices]

for i in range(len(start)):

v.append(self._data[start[i]: end[i]].astype(np.int32))

# Form batch

batch = {k: [] for k in ["inputs", "labels", "position_ids"]}

seqlens_i, seqlens_o = [], []

for i, o in zip(raw["inst"], raw["resp"]):

# PrefixLM: last input token predict first output token, for compatibility with Causal LM frameworks (prefill-phase)

# Inputs

batch["inputs"].append(i)

batch["inputs"].append(o[:-1])

# Labels

batch["labels"].append(np.full(len(i) - 1, dtype=i.dtype, fill_value=IGNORE_LABEL_ID) if self.config.target_only else i[1:])

batch["labels"].append(o)

# Position IDs

batch["position_ids"].append(np.arange(len(i), dtype=np.int32))

batch["position_ids"].append(np.arange(len(i), len(i) + len(o) - 1, dtype=np.int32))

# Seqlens

seqlens_i.append(len(i))

seqlens_o.append(len(o) - 1)

# Concat

batch = {k: np.concatenate(v, dtype=np.int32) for k, v in batch.items()}

# pad to fixed len

pad_len = self.config.batch_max_length - batch["inputs"].shape[0]

if pad_len > 0:

pad_values = {

"inputs": 0, # FIXME: Pad with an arbitary token.

"labels": IGNORE_LABEL_ID,

"position_ids": 0,

}

for k in pad_values.keys():

batch[k] = np.pad(batch[k], (0, pad_len), mode="constant", constant_values=pad_values[k])

# Compute cu_seqlens

seqlen_i, seqlen_o = np.array(seqlens_i, dtype=np.int32), np.array(seqlens_o, dtype=np.int32)

tensors, scalars = compute_aux_seq_tensors_scalars(seqlen_i, seqlen_o, self.config.batch_max_length)

# to tensor

return {k: torch.from_numpy(v) for k, v in (batch | tensors).items()}, scalars

def __iter__(self):

worker_info = get_worker_info()

assert worker_info is None or worker_info.num_workers == 1

# TODO: Feature (Low Priority): Multithreaded data loading

self._load_dataset_before_epoch_begin()

assert self._sampler is not None

for indices in self._sampler.iter():