docs: fix storage backends (2->3) and per-field length filter
- Storage Backends: add JsonlStore (lazy on-the-fly tokenization) - Describe memory-efficiency tradeoff across three backends - Fix length filter: per-text-field skip, not per-sample global - Abstract: HDF5/mmap -> tiered storage backends - Rename metric figures (singular) and update references
This commit is contained in:
Binary file not shown.
|
After Width: | Height: | Size: 92 KiB |
Binary file not shown.
|
Before Width: | Height: | Size: 166 KiB |
Binary file not shown.
|
Before Width: | Height: | Size: 128 KiB |
Binary file not shown.
|
After Width: | Height: | Size: 68 KiB |
Binary file not shown.
|
After Width: | Height: | Size: 85 KiB |
Binary file not shown.
|
Before Width: | Height: | Size: 172 KiB |
@@ -30,8 +30,8 @@
|
|||||||
\begin{abstract}
|
\begin{abstract}
|
||||||
We present {\sc AstrAI}, an open-source framework for end-to-end training
|
We present {\sc AstrAI}, an open-source framework for end-to-end training
|
||||||
of a 1.2B-parameter Transformer on $\sim$20B tokens. The pipeline covers
|
of a 1.2B-parameter Transformer on $\sim$20B tokens. The pipeline covers
|
||||||
JSON-driven BBPE preprocessing with multi-strategy packing, HDF5/mmap
|
JSON-driven BBPE preprocessing with multi-strategy packing, tiered
|
||||||
storage backends, and a companion SFT pipeline ({\sc Alembic}) with MinHash
|
storage backends (in-memory HDF5, memory-mapped binary, and lazy JSONL), and a companion SFT pipeline ({\sc Alembic}) with MinHash
|
||||||
deduplication and LLM-as-Judge scoring. The 24-layer decoder uses GQA,
|
deduplication and LLM-as-Judge scoring. The 24-layer decoder uses GQA,
|
||||||
SwiGLU, RoPE, and RMSNorm, trained with a hybrid Muon/AdamW optimizer and
|
SwiGLU, RoPE, and RMSNorm, trained with a hybrid Muon/AdamW optimizer and
|
||||||
WSD (Warmup--Stable--Decay) scheduling under DDP/FSDP. A BF16 stability analysis shows that
|
WSD (Warmup--Stable--Decay) scheduling under DDP/FSDP. A BF16 stability analysis shows that
|
||||||
@@ -91,17 +91,30 @@ via a JSON specification that defines:
|
|||||||
\texttt{.bin} shards, auto-split at 100M tokens per shard.
|
\texttt{.bin} shards, auto-split at 100M tokens per shard.
|
||||||
\end{itemize}
|
\end{itemize}
|
||||||
|
|
||||||
Samples shorter than 50~chars or longer than 2M~chars are filtered out.
|
In text-mode sections, individual fields shorter than 50~chars or
|
||||||
|
longer than 2M~chars are skipped during tokenization.
|
||||||
|
|
||||||
\subsection{Storage Backends}
|
\subsection{Storage Backends}
|
||||||
|
|
||||||
Two storage backends serve the DataLoader:
|
Three storage backends serve the DataLoader, trading off memory
|
||||||
|
footprint against access speed for datasets ranging from fine-tuning
|
||||||
|
scale to TB-level pretraining:
|
||||||
|
|
||||||
\begin{itemize}[nosep]
|
\begin{itemize}[nosep]
|
||||||
\item \textbf{H5Store}: HDF5-based, memory-loaded with shared-memory support
|
\item \textbf{H5Store}: HDF5-based, fully loaded into RAM at init
|
||||||
for multi-worker access.
|
with \texttt{share\_memory\_()} for cross-worker sharing.
|
||||||
\item \textbf{MmapStore}: Zero-copy memory-mapped \texttt{.bin} files shared
|
Fastest random access; requires the dataset to fit in memory.
|
||||||
via OS page cache.
|
\item \textbf{MmapStore}: Zero-copy memory-mapped \texttt{.bin} files
|
||||||
|
via \texttt{np.memmap}. Data stays on disk, managed by the OS
|
||||||
|
page cache; multiple workers share physical pages without
|
||||||
|
duplication. Suited for large pretraining corpora that exceed
|
||||||
|
RAM.
|
||||||
|
\item \textbf{JsonlStore}: Reads raw \texttt{.jsonl} directly. Lazy
|
||||||
|
mode (\texttt{processor=fn}) keeps only raw text records in
|
||||||
|
memory and defers per-sample tokenization to
|
||||||
|
\texttt{fetch\_record()}, avoiding a pre-tokenized copy
|
||||||
|
entirely---used by DPO/GRPO for on-the-fly training from
|
||||||
|
source files.
|
||||||
\end{itemize}
|
\end{itemize}
|
||||||
|
|
||||||
A resumable distributed sampler provides seed-based shuffle with
|
A resumable distributed sampler provides seed-based shuffle with
|
||||||
@@ -305,17 +318,17 @@ Figure~\ref{fig:ckpt_comparison} compares four configurations. The left panel sh
|
|||||||
|
|
||||||
\begin{figure}[H]
|
\begin{figure}[H]
|
||||||
\centering
|
\centering
|
||||||
\includegraphics[width=0.85\linewidth]{data/muon_pt.png}
|
\includegraphics[width=0.85\linewidth]{data/pt_metric.png}
|
||||||
\caption{Muon optimizer training dynamics: training loss and learning
|
\caption{Muon optimizer training dynamics: training loss and learning
|
||||||
rate schedule for the hybrid Muon/AdamW configuration across $\sim$25B
|
rate schedule for the hybrid Muon/AdamW configuration across $\sim$25B
|
||||||
tokens. The Muon optimizer is applied to all 2D weight matrices
|
tokens. The Muon optimizer is applied to all 2D weight matrices
|
||||||
(attention projections and FFN layers), while AdamW handles 1D
|
(attention projections and FFN layers), while AdamW handles 1D
|
||||||
parameters (embeddings, biases, and normalization scales). The WSD
|
parameters (embeddings, biases, and normalization scales). The WSD
|
||||||
schedule (warmup, stable phase, and final decay) is visible in the lower panel.}
|
schedule (warmup, stable phase, and final decay) is visible in the lower panel.}
|
||||||
\label{fig:muon_pt}
|
\label{fig:pt_metric}
|
||||||
\end{figure}
|
\end{figure}
|
||||||
|
|
||||||
Figure~\ref{fig:muon_pt} shows the extended training dynamics of the
|
Figure~\ref{fig:pt_metric} shows the extended training dynamics of the
|
||||||
Muon optimizer configuration over $\sim$25B tokens. The loss curve
|
Muon optimizer configuration over $\sim$25B tokens. The loss curve
|
||||||
exhibits the expected power-law decay in the early phase (0--5B tokens),
|
exhibits the expected power-law decay in the early phase (0--5B tokens),
|
||||||
followed by a gradual plateau as the model approaches convergence on
|
followed by a gradual plateau as the model approaches convergence on
|
||||||
@@ -325,7 +338,7 @@ training, ensuring stable weight updates in the final phase.
|
|||||||
|
|
||||||
\begin{figure}[H]
|
\begin{figure}[H]
|
||||||
\centering
|
\centering
|
||||||
\includegraphics[width=0.85\linewidth]{data/sft_metrics.png}
|
\includegraphics[width=0.85\linewidth]{data/sft_metric.png}
|
||||||
\caption{SFT training metrics over 1{,}000 fine-tuning steps on a
|
\caption{SFT training metrics over 1{,}000 fine-tuning steps on a
|
||||||
mixed Chinese--English instruction dataset: training loss, learning
|
mixed Chinese--English instruction dataset: training loss, learning
|
||||||
rate, and gradient norm. The loss drops rapidly in the first 200
|
rate, and gradient norm. The loss drops rapidly in the first 200
|
||||||
@@ -333,10 +346,10 @@ steps and then enters a slower decay phase. The learning rate
|
|||||||
follows a WSD schedule (2\% warmup, stable, decay). Gradient norms stabilize
|
follows a WSD schedule (2\% warmup, stable, decay). Gradient norms stabilize
|
||||||
after approximately 300 steps, indicating that the fine-tuning
|
after approximately 300 steps, indicating that the fine-tuning
|
||||||
process has reached a stable optimization regime.}
|
process has reached a stable optimization regime.}
|
||||||
\label{fig:sft_metrics}
|
\label{fig:sft_metric}
|
||||||
\end{figure}
|
\end{figure}
|
||||||
|
|
||||||
Figure~\ref{fig:sft_metrics} shows the supervised fine-tuning metrics
|
Figure~\ref{fig:sft_metric} shows the supervised fine-tuning metrics
|
||||||
for the 1K-step SFT checkpoint used in the IFD analysis
|
for the 1K-step SFT checkpoint used in the IFD analysis
|
||||||
(Appendix~\ref{app:ifd}). The training loss on the mixed
|
(Appendix~\ref{app:ifd}). The training loss on the mixed
|
||||||
Chinese--English instruction dataset decreases from $\sim$2.5 to
|
Chinese--English instruction dataset decreases from $\sim$2.5 to
|
||||||
@@ -368,18 +381,18 @@ cosine schedule with short linear warmup. The shorter ~3{,}000-step
|
|||||||
alignment run does not benefit from an extended stable phase; instead,
|
alignment run does not benefit from an extended stable phase; instead,
|
||||||
cosine decay lowers the learning rate steadily, which discourages
|
cosine decay lowers the learning rate steadily, which discourages
|
||||||
over-optimisation away from the reference distribution and matches the
|
over-optimisation away from the reference distribution and matches the
|
||||||
observed convergence pattern in Figure~\ref{fig:dpo_metrics}.
|
observed convergence pattern in Figure~\ref{fig:dpo_metric}.
|
||||||
|
|
||||||
\begin{figure}[H]
|
\begin{figure}[H]
|
||||||
\centering
|
\centering
|
||||||
\includegraphics[width=0.95\linewidth]{data/dpo_metrics.png}
|
\includegraphics[width=0.95\linewidth]{data/dpo_metric.png}
|
||||||
\caption{DPO training metrics over $\sim$3{,}000 alignment steps:
|
\caption{DPO training metrics over $\sim$3{,}000 alignment steps:
|
||||||
preference loss (raw and 100-step moving average), learning rate
|
preference loss (raw and 100-step moving average), learning rate
|
||||||
schedule, and gradient norm.}
|
schedule, and gradient norm.}
|
||||||
\label{fig:dpo_metrics}
|
\label{fig:dpo_metric}
|
||||||
\end{figure}
|
\end{figure}
|
||||||
|
|
||||||
Figure~\ref{fig:dpo_metrics} summarises the DPO training dynamics.
|
Figure~\ref{fig:dpo_metric} summarises the DPO training dynamics.
|
||||||
The raw preference loss (left panel) starts near $0.69$ and is visibly
|
The raw preference loss (left panel) starts near $0.69$ and is visibly
|
||||||
noisy, a hallmark of pairwise preference sampling. The 100-step moving
|
noisy, a hallmark of pairwise preference sampling. The 100-step moving
|
||||||
average reveals a steady downward trend that reaches a minimum of
|
average reveals a steady downward trend that reaches a minimum of
|
||||||
|
|||||||
Reference in New Issue
Block a user