128 lines
6.0 KiB
C
128 lines
6.0 KiB
C
/* SPDX-License-Identifier: AGPL-3.0-or-later */
|
|
/* Copyright (C) 2026 David (Corey) Herrera */
|
|
#ifndef FOUNDRY_BYTE_DATASET_H
|
|
#define FOUNDRY_BYTE_DATASET_H
|
|
|
|
#include <stdint.h>
|
|
#include <stddef.h>
|
|
|
|
#ifdef __cplusplus
|
|
extern "C" {
|
|
#endif
|
|
|
|
/*
|
|
* fd_byte_dataset — byte-level, mmap'd next-byte-prediction data loader.
|
|
*
|
|
* Built for the RWKV-x SR4-vs-AdamW scale test. The corpus is a raw
|
|
* byte file (enwik8); a BYTE IS A TOKEN, so the vocabulary is exactly 256 and
|
|
* there is no tokenizer / no BPE. The file is mmap'd PROT_READ (no full-file
|
|
* load; pages fault in on demand), and windows are sampled directly out of the
|
|
* mapping.
|
|
*
|
|
* Sampling (next-byte prediction), for a window whose start byte is p:
|
|
* input [j] = bytes[p + j] for j in [0, seq_len)
|
|
* target[j] = bytes[p + j + 1] for j in [0, seq_len) (input shifted +1)
|
|
* so a window consumes seq_len+1 contiguous bytes and the last valid start is
|
|
* (length - seq_len - 1). The number of valid window starts is:
|
|
* n_windows = length - seq_len
|
|
*
|
|
* DETERMINISM (mandatory for the reproducible A/B). The order in which window
|
|
* start-positions are visited is a seeded permutation derived ONLY from
|
|
* (seed, epoch); which windows land at a given (step, rank) is derived from
|
|
* (step, rank, num_ranks, batch). Therefore:
|
|
* same (seed, epoch, step, rank, num_ranks, batch) ⇒ byte-identical windows.
|
|
* The permutation is a cycle-walking Feistel network over the window-index
|
|
* domain (a true bijection on [0, n_windows) for any n_windows), so no
|
|
* O(n_windows) shuffle array is materialised. seed == FD_BYTE_SEED_UNSET (0) is
|
|
* rejected (fail-loud) so an unseeded run can never silently be nondeterministic.
|
|
*/
|
|
|
|
/* Vocabulary is exactly 256: every byte value 0..255 is its own token id. */
|
|
#define FD_BYTE_VOCAB 256
|
|
|
|
/* Reserved sentinel: an explicit seed of 0 means "unseeded" and is rejected. */
|
|
#define FD_BYTE_SEED_UNSET 0ULL
|
|
|
|
/* Standard enwik8 split (bytes). first 90M train, next 5M val, last 5M test. */
|
|
#define FD_ENWIK8_SIZE 100000000UL
|
|
#define FD_ENWIK8_TRAIN_OFF 0UL
|
|
#define FD_ENWIK8_TRAIN_LEN 90000000UL
|
|
#define FD_ENWIK8_VAL_OFF 90000000UL
|
|
#define FD_ENWIK8_VAL_LEN 5000000UL
|
|
#define FD_ENWIK8_TEST_OFF 95000000UL
|
|
#define FD_ENWIK8_TEST_LEN 5000000UL
|
|
|
|
/* mmap'd byte corpus (the whole file, read-only). */
|
|
typedef struct {
|
|
int fd; /* open file descriptor (kept for munmap accounting) */
|
|
const uint8_t *data; /* mmap base, PROT_READ */
|
|
size_t size; /* file size in bytes */
|
|
} fd_byte_corpus;
|
|
|
|
/* A view onto one split of the corpus: a contiguous [offset, offset+length)
|
|
* byte range plus the window length. Never reads outside its own range. */
|
|
typedef struct {
|
|
const uint8_t *base; /* corpus->data + offset (window byte 0 lives here) */
|
|
size_t offset; /* byte offset of this split within the corpus */
|
|
size_t length; /* split length in bytes */
|
|
size_t seq_len; /* window length (tokens per input/target row) */
|
|
size_t n_windows; /* # valid window starts = length - seq_len */
|
|
} fd_byte_split;
|
|
|
|
/* ------------------------------------------------------------------ */
|
|
/* Corpus (mmap) */
|
|
/* ------------------------------------------------------------------ */
|
|
|
|
/* mmap the byte file read-only. Returns 0 on success, -1 on failure
|
|
* (errno set; nothing allocated). */
|
|
int fd_byte_corpus_open(fd_byte_corpus *c, const char *path);
|
|
|
|
/* munmap + close. Safe on a zeroed / already-closed struct. */
|
|
void fd_byte_corpus_close(fd_byte_corpus *c);
|
|
|
|
/* ------------------------------------------------------------------ */
|
|
/* Splits */
|
|
/* ------------------------------------------------------------------ */
|
|
|
|
/* Initialise an arbitrary split view. Fails (-1) if [offset, offset+length)
|
|
* exceeds the corpus, or if length <= seq_len (no full window fits). */
|
|
int fd_byte_split_init(fd_byte_split *s, const fd_byte_corpus *c,
|
|
size_t offset, size_t length, size_t seq_len);
|
|
|
|
/* Convenience: the three standard enwik8 splits from the offsets above.
|
|
* Each requires the corpus to be exactly FD_ENWIK8_SIZE bytes. Returns 0/-1. */
|
|
int fd_byte_split_enwik8_train(fd_byte_split *s, const fd_byte_corpus *c, size_t seq_len);
|
|
int fd_byte_split_enwik8_val (fd_byte_split *s, const fd_byte_corpus *c, size_t seq_len);
|
|
int fd_byte_split_enwik8_test (fd_byte_split *s, const fd_byte_corpus *c, size_t seq_len);
|
|
|
|
/* ------------------------------------------------------------------ */
|
|
/* Deterministic window sampling */
|
|
/* ------------------------------------------------------------------ */
|
|
|
|
/* Map an ordinal position in the per-epoch stream to a window START BYTE
|
|
* (relative to the split base). This is the seeded permutation applied to the
|
|
* ordinal (mod n_windows). Deterministic in (seed, epoch, ordinal). Aborts if
|
|
* seed == FD_BYTE_SEED_UNSET. Returned p satisfies 0 <= p < n_windows, so the
|
|
* window [p, p+seq_len] stays strictly inside the split (no cross-split read). */
|
|
size_t fd_byte_window_pos(const fd_byte_split *s,
|
|
uint64_t seed, uint64_t epoch, uint64_t ordinal);
|
|
|
|
/* Fill one (batch, seq_len) minibatch of int token ids for this split.
|
|
* inputs / targets : caller-allocated, each batch*seq_len ints.
|
|
* inputs [b*seq_len + j] = byte[p_b + j]
|
|
* targets[b*seq_len + j] = byte[p_b + j + 1]
|
|
* where p_b is the window start for global ordinal
|
|
* ordinal = (step * num_ranks + rank) * batch + b
|
|
* dealt through the (seed, epoch) permutation. Every id is in [0, 256).
|
|
* Returns 0 on success, -1 on bad args. Aborts (fail-loud) if unseeded. */
|
|
int fd_byte_get_batch(const fd_byte_split *s,
|
|
uint64_t seed, uint64_t epoch, uint64_t step,
|
|
uint32_t rank, uint32_t num_ranks, uint32_t batch,
|
|
int *inputs, int *targets);
|
|
|
|
#ifdef __cplusplus
|
|
}
|
|
#endif
|
|
|
|
#endif /* FOUNDRY_BYTE_DATASET_H */
|