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:
2026-07-20 11:27:39 +08:00
parent da0f536526
commit d93ff48320
7 changed files with 31 additions and 18 deletions
Binary file not shown.

After

Width:  |  Height:  |  Size: 92 KiB

Binary file not shown.

Before

Width:  |  Height:  |  Size: 166 KiB

BIN
View File
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

+31 -18
View File
@@ -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