diff --git a/data/dpo_metric.png b/data/dpo_metric.png new file mode 100644 index 0000000..4fd4b33 Binary files /dev/null and b/data/dpo_metric.png differ diff --git a/data/dpo_metrics.png b/data/dpo_metrics.png deleted file mode 100644 index bdf18c1..0000000 Binary files a/data/dpo_metrics.png and /dev/null differ diff --git a/data/muon_pt.png b/data/muon_pt.png deleted file mode 100644 index 6c1b092..0000000 Binary files a/data/muon_pt.png and /dev/null differ diff --git a/data/pt_metric.png b/data/pt_metric.png new file mode 100644 index 0000000..e3b6e41 Binary files /dev/null and b/data/pt_metric.png differ diff --git a/data/sft_metric.png b/data/sft_metric.png new file mode 100644 index 0000000..2e4d4c0 Binary files /dev/null and b/data/sft_metric.png differ diff --git a/data/sft_metrics.png b/data/sft_metrics.png deleted file mode 100644 index c16f7ca..0000000 Binary files a/data/sft_metrics.png and /dev/null differ diff --git a/main.tex b/main.tex index dc1be88..cc0cc9c 100644 --- a/main.tex +++ b/main.tex @@ -30,8 +30,8 @@ \begin{abstract} 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 -JSON-driven BBPE preprocessing with multi-strategy packing, HDF5/mmap -storage backends, and a companion SFT pipeline ({\sc Alembic}) with MinHash +JSON-driven BBPE preprocessing with multi-strategy packing, tiered +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, 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 @@ -91,17 +91,30 @@ via a JSON specification that defines: \texttt{.bin} shards, auto-split at 100M tokens per shard. \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} -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] - \item \textbf{H5Store}: HDF5-based, memory-loaded with shared-memory support - for multi-worker access. - \item \textbf{MmapStore}: Zero-copy memory-mapped \texttt{.bin} files shared - via OS page cache. + \item \textbf{H5Store}: HDF5-based, fully loaded into RAM at init + with \texttt{share\_memory\_()} for cross-worker sharing. + Fastest random access; requires the dataset to fit in memory. + \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} 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] \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 rate schedule for the hybrid Muon/AdamW configuration across $\sim$25B tokens. The Muon optimizer is applied to all 2D weight matrices (attention projections and FFN layers), while AdamW handles 1D parameters (embeddings, biases, and normalization scales). The WSD schedule (warmup, stable phase, and final decay) is visible in the lower panel.} -\label{fig:muon_pt} +\label{fig:pt_metric} \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 exhibits the expected power-law decay in the early phase (0--5B tokens), 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] \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 mixed Chinese--English instruction dataset: training loss, learning 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 after approximately 300 steps, indicating that the fine-tuning process has reached a stable optimization regime.} -\label{fig:sft_metrics} +\label{fig:sft_metric} \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 (Appendix~\ref{app:ifd}). The training loss on the mixed 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, cosine decay lowers the learning rate steadily, which discourages 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] \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: preference loss (raw and 100-step moving average), learning rate schedule, and gradient norm.} -\label{fig:dpo_metrics} +\label{fig:dpo_metric} \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 noisy, a hallmark of pairwise preference sampling. The 100-step moving average reveals a steady downward trend that reaches a minimum of