# masworks/ml-agent architecture

> Architecture diagram of the GitHub repository masworks/ml-agent, generated by GitDiagram from its file tree, README and a sample of source files (last updated 2026-05-13). It is an AI-made overview: check details against the source.

- Interactive diagram: https://gitdiagram.com/masworks/ml-agent
- Repository: https://github.com/masworks/ml-agent

## Overview

1. **What this repository is**
   This repo is an ML-agent training and evaluation project whose released code is primarily the **exploration-enriched fine-tuning** stage. The codebase is built on top of **LLaMA-Factory**, so the architecture is centered on LLM fine-tuning, data preprocessing, inference/chat serving, and evaluation. The README explicitly says the **step-wise RL stage lives in VeRL**, not here.

2. **Top-level execution paths**
   The main user entry points are:
   - `src/train.py` for training runs
   - `src/webui.py` for the interactive UI
   - `src/api.py` for API serving
   - `scripts/train_step1.sh` and `scripts/eval.sh` as repo-level orchestration scripts
   The `examples/` directory contains ready-to-run configs for SFT, LoRA, QLoRA, inference, merge, and accelerator-specific setups.

3. **Configuration and experiment management**
   The repo is config-driven. Hyperparameters are organized under `src/llamafactory/hparams/`, with separate argument groups for model, data, training, generation, and evaluation. The `examples/` tree shows the intended experiment matrix: full finetuning, LoRA, QLoRA, distributed training variants, and inference backends. This is the main boundary for defining an experiment before any model/data code runs.

4. **Data subsystem**
   The data path starts in `src/llamafactory/data/`:
   - dataset metadata and task definitions are in `data/dataset_info.json`
   - loaders/parsers/formatters/converters normalize raw datasets into training examples
   - processors split the behavior by objective: supervised, unsupervised, pretraining, pairwise, and feedback
   - `template.py` and `tool_utils.py` shape prompts, chat templates, and tool-style examples
   - `mm_plugin.py` adds multimodal handling for image/vision-language tasks
   In architectural terms, this layer converts heterogeneous benchmark/task data into a unified tokenized format consumed by the trainers.

5. **Model loading and adaptation**
   The `src/llamafactory/model/` layer is responsible for instantiating base models and applying training-time modifications. It includes:
   - `loader.py` for model construction
   - `adapter.py` for parameter-efficient tuning paths such as LoRA-style adapters
   - `patcher.py` and `model_utils/` for runtime model edits
   - support modules for quantization, RoPE, attention, checkpointing, embedding, MoE, visual components, and long-context behavior
   This is the main boundary between the generic Hugging Face model stack and the project-specific training strategies.

6. **Training workflows**
   Training is split by objective under `src/llamafactory/train/`:
   - `sft/` for supervised fine-tuning
   - `pt/` for pretraining
   - `rm/` for reward modeling
   - `dpo/` for preference optimization
   - `kto/` and `ppo/` for additional alignment/RL-style methods
   Each mode has a `workflow.py` that wires arguments, data, model, and trainer together, plus a `trainer.py` that encapsulates the actual optimization loop. Shared logic lives in `trainer_utils.py`, `callbacks.py`, and `tuner.py`.

7. **Evaluation subsystem**
   Evaluation is handled separately from training in `src/llamafactory/eval/`, with an evaluator and templates for task-specific scoring. The standalone `evaluation/` directory contains benchmark-specific scripts and mappings for:
   - MMLU
   - CMMLU
   - CEval
   This indicates a clean separation between model training and benchmark evaluation, with evaluation scripts adapting model outputs to benchmark scoring rules.

8. **Inference, chat, and serving**
   The runtime inference boundary is in `src/llamafactory/chat/` and `src/llamafactory/api/`:
   - `chat_model.py` and engine implementations wrap either Hugging Face or **vLLM**
   - `base_engine.py`, `hf_engine.py`, and `vllm_engine.py` select execution backends
   - `api/app.py`, `api/chat.py`, and `api/protocol.py` expose serving endpoints and request/response schemas
   This layer is for post-training use of the model, not for optimization. The repo also includes `scripts/vllm_infer.py` for direct vLLM-based inference.

9. **Web UI as an orchestration layer**
   The `src/llamafactory/webui/` package is a Gradio-style interactive front end that wraps training, inference, evaluation, data management, and export flows. Components are split by function:
   - chat
   - data
   - eval
   - infer
   - train
   - export
   Supporting modules like `runner.py`, `engine.py`, and `manager.py` coordinate background execution and state management. Architecturally, the UI does not own model logic; it delegates to the same underlying training/inference subsystems.

10. **Scripts and operational tooling**
   The repo includes utility scripts for:
   - checkpoint conversion (`scripts/convert_ckpt/`)
   - initializing specific adapter methods such as PiSSA and LoFTQ
   - statistics and diagnostics (`scripts/stat_utils/`)
   - API examples and multimodal/tool-call demos (`scripts/api_example/`)
   These scripts are operational glue around the core training stack and help move models between formats or inspect training behavior.

11. **External infrastructure and execution backends**
   The architecture assumes integration with the broader Hugging Face ecosystem and distributed training backends. The examples show support for:
   - **Accelerate**
   - **DeepSpeed**
   - **FSDP**
   - **vLLM**
   There are also specialized QLoRA/quantization paths, which affect memory layout and backend selection. These are important architectural boundaries because they change how the same trainer executes on different hardware budgets.

12. **Repository boundary relative to the ML-Agent paper**
   For the ML-Agent project described in the README, this repo covers the **fine-tuning foundation** and evaluation scaffolding, but not the full agentic training stack. The held-in/held-out ML task framing is reflected in the data/evaluation organization, while the online RL stage is intentionally externalized. If you are drawing the system diagram, treat this repository as the **model adaptation and benchmark-facing half** of the overall ML-Agent pipeline.

Main entry points: [`LLaMA-Factory/src/llamafactory/data`](https://github.com/masworks/ml-agent/tree/main/LLaMA-Factory/src/llamafactory/data) (data pipeline), [`LLaMA-Factory/src/llamafactory/webui`](https://github.com/masworks/ml-agent/tree/main/LLaMA-Factory/src/llamafactory/webui) (web UI), [`LLaMA-Factory/src/llamafactory/chat`](https://github.com/masworks/ml-agent/tree/main/LLaMA-Factory/src/llamafactory/chat) (chat engine), [`LLaMA-Factory/src/llamafactory/hparams`](https://github.com/masworks/ml-agent/tree/main/LLaMA-Factory/src/llamafactory/hparams) (hparams)

## Diagram (Mermaid)

```mermaid
flowchart TD

subgraph group_entry["Entry points"]
  node_src_train["train.py<br/>entrypoint<br/>[train.py]"]
  node_src_webui["webui.py<br/>entrypoint<br/>[webui.py]"]
  node_src_api["api.py<br/>entrypoint<br/>[api.py]"]
  node_cli["cli<br/>command layer<br/>[cli.py]"]
end

subgraph group_core["Core stack"]
  node_hparams["hparams<br/>config schema"]
  node_dataset_info["dataset info<br/>dataset registry<br/>[dataset_info.json]"]
  node_data_pipeline["data pipeline<br/>data prep"]
  node_data_templates["templates<br/>prompt shaping<br/>[template.py]"]
  node_mm_plugin["mm plugin<br/>multimodal adapter<br/>[mm_plugin.py]"]
  node_model_loader["model loader<br/>model init<br/>[loader.py]"]
  node_model_adapter["adapters<br/>tuning hooks<br/>[adapter.py]"]
  node_model_utils["model utils<br/>runtime mods"]
  node_train_sft["SFT<br/>workflow<br/>[workflow.py]"]
  node_train_align["alignment<br/>workflows"]
  node_train_shared["trainer utils<br/>shared training<br/>[trainer_utils.py]"]
  node_eval_core["evaluator<br/>evaluation core<br/>[evaluator.py]"]
end

subgraph group_runtime["Runtime"]
  node_chat_runtime["chat engine<br/>inference runtime"]
  node_api_layer["API layer<br/>serving"]
  node_webui["web UI<br/>orchestration UI"]
end

subgraph group_ops["Ops and benchmarks"]
  node_examples["examples<br/>experiment matrix"]
  node_benchmarks["benchmarks<br/>benchmark suite"]
  node_scripts["scripts<br/>ops tooling"]
end

node_src_train -->|"dispatches"| node_cli
node_src_webui -->|"launches"| node_webui
node_src_api -->|"launches"| node_api_layer
node_cli -->|"parses"| node_hparams
node_hparams -.->|"instantiated by"| node_examples
node_hparams -->|"configures"| node_data_pipeline
node_dataset_info -->|"drives"| node_data_pipeline
node_data_pipeline -->|"formats with"| node_data_templates
node_data_pipeline -->|"extends"| node_mm_plugin
node_data_pipeline -->|"feeds"| node_train_sft
node_model_loader -->|"applies"| node_model_adapter
node_model_adapter -->|"uses"| node_model_utils
node_hparams -->|"configures"| node_model_loader
node_train_sft -->|"shares"| node_train_shared
node_train_align -->|"shares"| node_train_shared
node_train_align -->|"runs on"| node_model_loader
node_train_align -->|"consumes"| node_data_pipeline
node_eval_core -->|"adapts to"| node_benchmarks
node_chat_runtime -->|"loads"| node_model_loader
node_chat_runtime -->|"selects"| node_model_utils
node_api_layer -->|"wraps"| node_chat_runtime
node_webui -->|"orchestrates"| node_train_align
node_webui -->|"orchestrates"| node_chat_runtime
node_webui -->|"orchestrates"| node_eval_core
node_webui -->|"orchestrates"| node_data_pipeline
node_scripts -.->|"supports"| node_model_adapter
node_scripts -.->|"supports"| node_chat_runtime

click node_src_train "https://github.com/masworks/ml-agent/blob/main/LLaMA-Factory/src/train.py"
click node_src_webui "https://github.com/masworks/ml-agent/blob/main/LLaMA-Factory/src/webui.py"
click node_src_api "https://github.com/masworks/ml-agent/blob/main/LLaMA-Factory/src/api.py"
click node_cli "https://github.com/masworks/ml-agent/blob/main/LLaMA-Factory/src/llamafactory/cli.py"
click node_hparams "https://github.com/masworks/ml-agent/tree/main/LLaMA-Factory/src/llamafactory/hparams"
click node_examples "https://github.com/masworks/ml-agent/tree/main/LLaMA-Factory/examples"
click node_dataset_info "https://github.com/masworks/ml-agent/blob/main/LLaMA-Factory/data/dataset_info.json"
click node_data_pipeline "https://github.com/masworks/ml-agent/tree/main/LLaMA-Factory/src/llamafactory/data"
click node_data_templates "https://github.com/masworks/ml-agent/blob/main/LLaMA-Factory/src/llamafactory/data/template.py"
click node_mm_plugin "https://github.com/masworks/ml-agent/blob/main/LLaMA-Factory/src/llamafactory/data/mm_plugin.py"
click node_model_loader "https://github.com/masworks/ml-agent/blob/main/LLaMA-Factory/src/llamafactory/model/loader.py"
click node_model_adapter "https://github.com/masworks/ml-agent/blob/main/LLaMA-Factory/src/llamafactory/model/adapter.py"
click node_model_utils "https://github.com/masworks/ml-agent/tree/main/LLaMA-Factory/src/llamafactory/model/model_utils"
click node_train_sft "https://github.com/masworks/ml-agent/blob/main/LLaMA-Factory/src/llamafactory/train/sft/workflow.py"
click node_train_align "https://github.com/masworks/ml-agent/tree/main/LLaMA-Factory/src/llamafactory/train"
click node_train_shared "https://github.com/masworks/ml-agent/blob/main/LLaMA-Factory/src/llamafactory/train/trainer_utils.py"
click node_eval_core "https://github.com/masworks/ml-agent/blob/main/LLaMA-Factory/src/llamafactory/eval/evaluator.py"
click node_benchmarks "https://github.com/masworks/ml-agent/tree/main/LLaMA-Factory/evaluation"
click node_chat_runtime "https://github.com/masworks/ml-agent/tree/main/LLaMA-Factory/src/llamafactory/chat"
click node_api_layer "https://github.com/masworks/ml-agent/tree/main/LLaMA-Factory/src/llamafactory/api"
click node_webui "https://github.com/masworks/ml-agent/tree/main/LLaMA-Factory/src/llamafactory/webui"
click node_scripts "https://github.com/masworks/ml-agent/tree/main/LLaMA-Factory/scripts"

classDef toneNeutral fill:#f8fafc,stroke:#334155,stroke-width:1.5px,color:#0f172a
classDef toneBlue fill:#dbeafe,stroke:#2563eb,stroke-width:1.5px,color:#172554
classDef toneAmber fill:#fef3c7,stroke:#d97706,stroke-width:1.5px,color:#78350f
classDef toneMint fill:#dcfce7,stroke:#16a34a,stroke-width:1.5px,color:#14532d
classDef toneRose fill:#ffe4e6,stroke:#e11d48,stroke-width:1.5px,color:#881337
classDef toneIndigo fill:#e0e7ff,stroke:#4f46e5,stroke-width:1.5px,color:#312e81
classDef toneTeal fill:#ccfbf1,stroke:#0f766e,stroke-width:1.5px,color:#134e4a
class node_src_train,node_src_webui,node_src_api,node_cli toneBlue
class node_hparams,node_dataset_info,node_data_pipeline,node_data_templates,node_mm_plugin,node_model_loader,node_model_adapter,node_model_utils,node_train_sft,node_train_align,node_train_shared,node_eval_core toneAmber
class node_chat_runtime,node_api_layer,node_webui toneMint
class node_examples,node_benchmarks,node_scripts toneRose
```

## Components

### Entry points

Primary ways to launch training, UI, and serving.

- [train.py](https://github.com/masworks/ml-agent/blob/main/LLaMA-Factory/src/train.py): entrypoint · Main CLI for training runs. · `LLaMA-Factory/src/train.py`
- [webui.py](https://github.com/masworks/ml-agent/blob/main/LLaMA-Factory/src/webui.py): entrypoint · Launches the interactive Gradio-style UI. · `LLaMA-Factory/src/webui.py`
- [api.py](https://github.com/masworks/ml-agent/blob/main/LLaMA-Factory/src/api.py): entrypoint · Starts the API serving process. · `LLaMA-Factory/src/api.py`
- [cli](https://github.com/masworks/ml-agent/blob/main/LLaMA-Factory/src/llamafactory/cli.py): command layer · Shared command-line dispatch and execution entry logic. · `LLaMA-Factory/src/llamafactory/cli.py`

### Core stack

Configuration, data prep, model adaptation, and training workflows.

- [hparams](https://github.com/masworks/ml-agent/tree/main/LLaMA-Factory/src/llamafactory/hparams): config schema · Argument groups and parser for model, data, training, generation, and evaluation. · `LLaMA-Factory/src/llamafactory/hparams`
- [dataset info](https://github.com/masworks/ml-agent/blob/main/LLaMA-Factory/data/dataset_info.json): dataset registry · Dataset metadata and task mapping registry. · `LLaMA-Factory/data/dataset_info.json`
- [data pipeline](https://github.com/masworks/ml-agent/tree/main/LLaMA-Factory/src/llamafactory/data): data prep · Loads, parses, formats, converts, and collates heterogeneous datasets. · `LLaMA-Factory/src/llamafactory/data`
- [templates](https://github.com/masworks/ml-agent/blob/main/LLaMA-Factory/src/llamafactory/data/template.py): prompt shaping · Chat templates, tool-use formatting, and multimodal prompt shaping. · `LLaMA-Factory/src/llamafactory/data/template.py`
- [mm plugin](https://github.com/masworks/ml-agent/blob/main/LLaMA-Factory/src/llamafactory/data/mm_plugin.py): multimodal adapter · Adds image and vision-language handling to the data path. · `LLaMA-Factory/src/llamafactory/data/mm_plugin.py`
- [model loader](https://github.com/masworks/ml-agent/blob/main/LLaMA-Factory/src/llamafactory/model/loader.py): model init · Constructs base models and applies runtime/model-specific setup. · `LLaMA-Factory/src/llamafactory/model/loader.py`
- [adapters](https://github.com/masworks/ml-agent/blob/main/LLaMA-Factory/src/llamafactory/model/adapter.py): tuning hooks · Applies LoRA-style adaptation and model patching. · `LLaMA-Factory/src/llamafactory/model/adapter.py`
- [model utils](https://github.com/masworks/ml-agent/tree/main/LLaMA-Factory/src/llamafactory/model/model_utils): runtime mods · Quantization, RoPE, attention, checkpointing, visual, and other model edits. · `LLaMA-Factory/src/llamafactory/model/model_utils`
- [SFT](https://github.com/masworks/ml-agent/blob/main/LLaMA-Factory/src/llamafactory/train/sft/workflow.py): workflow · Supervised fine-tuning workflow and trainer. · `LLaMA-Factory/src/llamafactory/train/sft/workflow.py`
- [alignment](https://github.com/masworks/ml-agent/tree/main/LLaMA-Factory/src/llamafactory/train): workflows · Preference and alignment workflows for PT, RM, DPO, KTO, and PPO. · `LLaMA-Factory/src/llamafactory/train`
- [trainer utils](https://github.com/masworks/ml-agent/blob/main/LLaMA-Factory/src/llamafactory/train/trainer_utils.py): shared training · Common trainer utilities, callbacks, and tuning helpers. · `LLaMA-Factory/src/llamafactory/train/trainer_utils.py`
- [evaluator](https://github.com/masworks/ml-agent/blob/main/LLaMA-Factory/src/llamafactory/eval/evaluator.py): evaluation core · General evaluation orchestration and scoring templates. · `LLaMA-Factory/src/llamafactory/eval/evaluator.py`

### Runtime

Inference, chat, and API serving backends.

- [chat engine](https://github.com/masworks/ml-agent/tree/main/LLaMA-Factory/src/llamafactory/chat): inference runtime · Chat model wrapper with Hugging Face and vLLM backend selection. · `LLaMA-Factory/src/llamafactory/chat`
- [API layer](https://github.com/masworks/ml-agent/tree/main/LLaMA-Factory/src/llamafactory/api): serving · Request protocol and HTTP handlers for model serving. · `LLaMA-Factory/src/llamafactory/api`
- [web UI](https://github.com/masworks/ml-agent/tree/main/LLaMA-Factory/src/llamafactory/webui): orchestration UI · Interactive front end that delegates to training, inference, evaluation, data, and export subsystems. · `LLaMA-Factory/src/llamafactory/webui`

### Ops and benchmarks

Repository scripts, examples, and external benchmark evaluation.

- [examples](https://github.com/masworks/ml-agent/tree/main/LLaMA-Factory/examples): experiment matrix · Ready-to-run configs for SFT, LoRA, QLoRA, inference, merge, and distributed setups. · `LLaMA-Factory/examples`
- [benchmarks](https://github.com/masworks/ml-agent/tree/main/LLaMA-Factory/evaluation): benchmark suite · Benchmark-specific scripts and mappings for MMLU, CMMLU, and CEval. · `LLaMA-Factory/evaluation`
- [scripts](https://github.com/masworks/ml-agent/tree/main/LLaMA-Factory/scripts): ops tooling · Utility scripts for checkpoint conversion, initialization, statistics, and direct inference. · `LLaMA-Factory/scripts`

## Connections

- train.py → cli: dispatches · Training entrypoint uses shared CLI dispatch.
- webui.py → web UI: launches · UI entrypoint starts the web interface.
- api.py → API layer: launches · API entrypoint starts the serving layer.
- cli → hparams: parses · CLI builds typed argument groups from config schema.
- hparams → examples: instantiated by · Example configs define concrete experiment variants.
- hparams → data pipeline: configures · Data-related arguments drive dataset loading and formatting.
- dataset info → data pipeline: drives · Dataset registry informs loader and task selection.
- data pipeline → templates: formats with · Unified examples are shaped by templates and tool prompts.
- data pipeline → mm plugin: extends · Multimodal tasks pass through the multimodal plugin.
- data pipeline → SFT: feeds · Prepared examples are consumed by training workflows.
- model loader → adapters: applies · Loaded base models receive adaptation and patching.
- adapters → model utils: uses · Adapters rely on runtime model utilities and patches.
- hparams → model loader: configures · Model and tuning arguments control model construction.
- SFT → trainer utils: shares · SFT uses common trainer utilities and callbacks.
- alignment → trainer utils: shares · Other training objectives reuse the same trainer support.
- alignment → model loader: runs on · All optimization modes build models through the loader boundary.
- alignment → data pipeline: consumes · Alignment workflows consume the normalized data pipeline.
- evaluator → benchmarks: adapts to · Benchmark scripts provide task-specific evaluation rules.
- chat engine → model loader: loads · Inference reuses the model loading boundary.
- chat engine → model utils: selects · Backend execution depends on model/runtime utilities.
- API layer → chat engine: wraps · HTTP serving is built on the chat runtime.
- web UI → alignment: orchestrates · UI delegates training actions to the training workflows.
- web UI → chat engine: orchestrates · UI delegates inference and chat actions to runtime backends.
- web UI → evaluator: orchestrates · UI triggers evaluation through the evaluation core.
- web UI → data pipeline: orchestrates · UI includes dataset management and preprocessing flows.
- scripts → adapters: supports · Utility scripts initialize or convert adapter-related checkpoints.
- scripts → chat engine: supports · Inference and API helper scripts target runtime backends.