Augmentoolkit Documentation

The atk-interface is a web-based Single Page Application (SPA) designed to provide a visual interface for interacting with a locally-running Augmentoolkit server. It allows users to manage configurations, handle input/output files, and monitor pipeline execution in real-time.

Key functional areas include:

• Navigation: Access to Configs, Inputs, Outputs, and Pipeline monitoring.
• Config Editing: Direct modification of configuration fields with path validation (warnings for non-existent input_dir or output_dir) and pipeline execution capabilities.
• Input/Output Management: A file browser interface for uploading, viewing, deleting, and downloading files. The output section also tracks task IDs for the current session.
• Pipeline Execution Monitoring: A dedicated view featuring a progress bar, real-time log streaming, and an expandable details box showing execution stats and the specific configuration used (with syntax highlighting).

Augmentoolkit is built on several key technical hypotheses and long-term goals:

Core Hypotheses

• Factual Recall: With proper continued pretraining and SFT, LLMs can be trained to understand and recall new factual domains.
• Specialization over Generalization: Specialized AI (in knowledge, opinion, or tone) is more useful and controllable for specific products than generalist models.
• Synthetic Data: High-quality synthetic data is a primary driver for making specialized models effective and is a pillar of open-source AI.
• Abstraction-driven Generation: Dataset generation can be standardized through highly-practical abstractions and unified interfaces.

Future Goals

• Enhanced Fact Learning: Improving the efficacy of factual domain injection.
• GRPO Reinforcement Learning: Implementing Group Relative Policy Optimization (GRPO) for arbitrary objectives (e.g., code or LLM-graded prompts).
• Advanced Training Styles: Developing capabilities for writing style and agentic training.

Augmentoolkit is a tool designed to create AI subject matter experts by generating high-quality datasets for Continued Pretraining and Supervised Fine Tuning (SFT). It uses LLM-powered pipelines to teach models new factual domains, including information the model has never encountered during its initial pretraining.

The Core Methodology

To teach a model new facts without causing it to overfit on document structure, Augmentoolkit creates many synthetic, varied representations of the target documents. The workflow typically involves:

1. Continued Pretraining: Performed for an extended period (often around 12 epochs) to ensure the loss is low and the documents are completely memorized.
2. Supervised Fine Tuning (SFT): Solidifies the knowledge by training on a mix of domain-specific conversational data and generic datasets (e.g., Capybara, Hermes, Bluemoon, Pippa, LMSys). This step uses a decent batch size and typically runs for 5 to 7 epochs.
3. Hallucination Resistance: The process incorporates domain-specific data to train the model to say it "does not know" about invented facts, correct faulty assumptions, and handle follow-up questions.

Key Capabilities

• Knowledge Expansion: Arbitrarily expands an LLM's knowledge cutoff to any chosen area.
• Belief/Opinion Control: Allows for the arbitrary selection of an LLM's opinions or beliefs about what is true.
• Cost-Effectiveness: Can be run on consumer hardware. A custom-trained 7B parameter dataset generation model is available to serve as the engine, allowing domain experts to be trained for $20 or less using open-source AI.

The RAG Data pipeline is designed to train LLMs on how to effectively utilize context provided by a Retrieval-Augmented Generation (RAG) system. It generates multi-turn conversational data where the model is presented with simulated RAG results (both relevant and deliberately irrelevant chunks) alongside a question derived from a source text chunk.

Core Goals

• Teach the model to synthesize answers using provided context when relevant.
• Teach the model to rely on internal knowledge when the context is irrelevant or unhelpful.
• Train the model to ignore misleading or noisy retrieved information.

Use Cases

• Improving RAG System Performance: Fine-tuning the generator component to better utilize retriever output.
• Training Robustness to Noisy Context: Reducing susceptibility to irrelevant information.
• Enhancing Contextual Reasoning: Teaching models to synthesize information from multiple snippets.
• Building Advanced QA Systems: Creating models that dynamically switch between provided context and internal knowledge.

The GRPO (Generative Reward Powered Optimization) pipeline is an experimental framework for fine-tuning LLMs using Reinforcement Learning (RL). Unlike standard data generation pipelines, this performs actual model training to align a base model toward desired behaviors using custom reward functions.

Key Characteristics

• Purpose: Optimizes models using reward signals (programmatic or LLM-based) rather than just supervised fine-tuning (SFT).
• Platform: Linux only. Requires a machine with a capable GPU (suitable for LoRA training) and vLLM installed in your virtual environment.
• Complexity: This is a beta/experimental feature. It does not support auto-resume. Customization (adding reward functions) requires Python coding.
• Workflow: Generates multiple candidate responses for prompts from input datasets $\rightarrow$ evaluates responses via reward functions $\rightarrow$ optimizes LoRA weights via GRPO.

Pipelines use configuration files to manage their arguments.

Configuration Rules:

• Argument Mapping: Pipeline arguments are mapped from config files as keyword arguments.
• super_config.yaml: This file controls which pipeline is executed with which config when running in CLI mode. It also contains a path_aliases section used by the API to track existing pipelines and their primary configs.
• Path Aliases: The path_aliases section allows the API to identify where pipelines exist and where their primary configs are located (useful for duplicating configs to create new pipelines).
• Convention: By convention, at least one config file for a pipeline should be stored in the same folder as the pipeline itself. If listed under path_aliases, this config will appear as a template in the interface.
• Placeholders: For fields without sensible defaults (like API keys), use the string "!!PLACEHOLDER!!". The interface recognizes this and flags the field to prevent users from forgetting to fill it in.

The Complete Factual Generation pipeline is the primary composition pipeline in Augmentoolkit. It is designed to orchestrate multiple specialized pipelines to generate a full suite of datasets required to create a domain-expert LLM.

By running this single pipeline, you automate the following workflow:

1. Document Processing: Cleans input documents (including PDFs) into clean text.
2. Pretraining Data Generation: Uses the Representation Variation pipeline to create varied representations and inferred facts.
3. SFT Data Generation: Creates diverse Supervised Fine-Tuning (SFT) datasets, including:
   • Factual QA via the Multi-Source Recall pipeline.
   • RAG-focused training data via the RAG Data pipeline.
   • Self-correction examples via the Correction Data pipeline.
4. Data Balancing: Combines domain-specific SFT data with generic SFT datasets, automatically balancing token counts.
5. Training Readiness: Produces Axolotl YAML configuration files for both continued pretraining and final SFT steps.

Essentially, you provide documents and configure API keys/models, and the pipeline outputs everything needed to train your expert model.