tocmo0nlord/axolotl
1
0
Fork 0
Code Issues Pull Requests Actions 2 Packages Projects Releases Wiki Activity
Files
79255ffdd7c52a79e0a75b279dce6b6335b47f78
axolotl/docs/custom_integrations.html
Quarto GHA Workflow Runner 79255ffdd7 Built site for gh-pages
2026-04-24 09:09:53 +00:00

4200 lines
318 KiB
HTML
Raw Blame History

This file contains ambiguous Unicode characters
This file contains Unicode characters that might be confused with other characters. If you think that this is intentional, you can safely ignore this warning. Use the Escape button to reveal them.
<!DOCTYPE html>
<html xmlns="http://www.w3.org/1999/xhtml" lang="en" xml:lang="en"><head>
<meta charset="utf-8">
<meta name="generator" content="quarto-1.9.37">
<meta name="viewport" content="width=device-width, initial-scale=1.0, user-scalable=yes">
<title>Custom Integrations – Axolotl</title>
<style>
/* Default styles provided by pandoc.
** See https://pandoc.org/MANUAL.html#variables-for-html for config info.
*/
code{white-space: pre-wrap;}
span.smallcaps{font-variant: small-caps;}
div.columns{display: flex; gap: min(4vw, 1.5em);}
div.column{flex: auto; overflow-x: auto;}
div.hanging-indent{margin-left: 1.5em; text-indent: -1.5em;}
ul.task-list{list-style: none;}
ul.task-list li input[type="checkbox"] {
width: 0.8em;
margin: 0 0.8em 0.2em -1em; /* quarto-specific, see https://github.com/quarto-dev/quarto-cli/issues/4556 */
vertical-align: middle;
}
/* CSS for syntax highlighting */
html { -webkit-text-size-adjust: 100%; }
pre > code.sourceCode { white-space: pre; position: relative; }
pre > code.sourceCode > span { display: inline-block; line-height: 1.25; }
pre > code.sourceCode > span:empty { height: 1.2em; }
.sourceCode { overflow: visible; }
code.sourceCode > span { color: inherit; text-decoration: inherit; }
div.sourceCode { margin: 1em 0; }
pre.sourceCode { margin: 0; }
@media screen {
div.sourceCode { overflow: auto; }
}
@media print {
pre > code.sourceCode { white-space: pre-wrap; }
pre > code.sourceCode > span { text-indent: -5em; padding-left: 5em; }
}
pre.numberSource code
{ counter-reset: source-line 0; }
pre.numberSource code > span
{ position: relative; left: -4em; counter-increment: source-line; }
pre.numberSource code > span > a:first-child::before
{ content: counter(source-line);
position: relative; left: -1em; text-align: right; vertical-align: baseline;
border: none; display: inline-block;
-webkit-touch-callout: none; -webkit-user-select: none;
-khtml-user-select: none; -moz-user-select: none;
-ms-user-select: none; user-select: none;
padding: 0 4px; width: 4em;
}
pre.numberSource { margin-left: 3em; padding-left: 4px; }
div.sourceCode
{ }
@media screen {
pre > code.sourceCode > span > a:first-child::before { text-decoration: underline; }
}
</style>
<script src="../site_libs/quarto-nav/quarto-nav.js"></script>
<script src="../site_libs/clipboard/clipboard.min.js"></script>
<script src="../site_libs/quarto-search/autocomplete.umd.js"></script>
<script src="../site_libs/quarto-search/fuse.min.js"></script>
<script src="../site_libs/quarto-search/quarto-search.js"></script>
<meta name="quarto:offset" content="../">
<link href="../favicon.jpg" rel="icon" type="image/jpeg">
<script src="../site_libs/quarto-html/quarto.js" type="module"></script>
<script src="../site_libs/quarto-html/tabsets/tabsets.js" type="module"></script>
<script src="../site_libs/quarto-html/popper.min.js"></script>
<script src="../site_libs/quarto-html/tippy.umd.min.js"></script>
<script src="../site_libs/quarto-html/anchor.min.js"></script>
<link href="../site_libs/quarto-html/tippy.css" rel="stylesheet">
<link href="../site_libs/quarto-html/quarto-syntax-highlighting-dark-d0ae9245876894da5ac7e18953ecc5cc.css" rel="stylesheet" id="quarto-text-highlighting-styles">
<script src="../site_libs/bootstrap/bootstrap.min.js"></script>
<link href="../site_libs/bootstrap/bootstrap-icons.css" rel="stylesheet">
<link href="../site_libs/bootstrap/bootstrap-ab6ebd6eb475c4578b58908bc314f719.min.css" rel="stylesheet" append-hash="true" id="quarto-bootstrap" data-mode="dark">
<script id="quarto-search-options" type="application/json">{
"location": "navbar",
"copy-button": false,
"collapse-after": 3,
"panel-placement": "end",
"type": "overlay",
"limit": 50,
"keyboard-shortcut": [
"f",
"/",
"s"
],
"show-item-context": false,
"language": {
"search-no-results-text": "No results",
"search-matching-documents-text": "matching documents",
"search-copy-link-title": "Copy link to search",
"search-hide-matches-text": "Hide additional matches",
"search-more-match-text": "more match in this document",
"search-more-matches-text": "more matches in this document",
"search-clear-button-title": "Clear",
"search-text-placeholder": "",
"search-detached-cancel-button-title": "Cancel",
"search-submit-button-title": "Submit",
"search-label": "Search"
}
}</script>
<script async="" src="https://www.googletagmanager.com/gtag/js?id=G-9KYCVJBNMQ"></script>
<script type="text/javascript">
window.dataLayer = window.dataLayer || [];
function gtag(){dataLayer.push(arguments);}
gtag('js', new Date());
gtag('config', 'G-9KYCVJBNMQ', { 'anonymize_ip': true});
</script>
<link rel="stylesheet" href="../styles.css">
</head>
<body class="nav-sidebar docked nav-fixed quarto-light">
<div id="quarto-search-results"></div>
<header id="quarto-header" class="headroom fixed-top">
<nav class="navbar navbar-expand " data-bs-theme="dark">
<div class="navbar-container container-fluid">
<div class="navbar-brand-container mx-auto">
<a href="../index.html" class="navbar-brand navbar-brand-logo">
<img src="../image/axolotl_logo_digital_white.svg" alt="" class="navbar-logo light-content">
<img src="../image/axolotl_logo_digital_white.svg" alt="" class="navbar-logo dark-content">
</a>
</div>
<div class="quarto-navbar-tools tools-wide tools-end">
<a href="https://twitter.com/axolotl_ai" title="" class="quarto-navigation-tool px-1" aria-label=""><i class="bi bi-twitter"></i></a>
<a href="https://github.com/axolotl-ai-cloud/axolotl/" title="" class="quarto-navigation-tool px-1" aria-label=""><i class="bi bi-github"></i></a>
<a href="https://discord.gg/7m9sfhzaf3" title="" class="quarto-navigation-tool px-1" aria-label=""><i class="bi bi-discord"></i></a>
</div>
<div id="quarto-search" class="" title="Search"></div>
</div> <!-- /container-fluid -->
</nav>
<nav class="quarto-secondary-nav">
<div class="container-fluid d-flex">
<button type="button" class="quarto-btn-toggle btn" data-bs-toggle="collapse" role="button" data-bs-target=".quarto-sidebar-collapse-item" aria-controls="quarto-sidebar" aria-expanded="false" aria-label="Toggle sidebar navigation" onclick="if (window.quartoToggleHeadroom) { window.quartoToggleHeadroom(); }">
<i class="bi bi-layout-text-sidebar-reverse"></i>
</button>
<nav class="quarto-page-breadcrumbs" aria-label="breadcrumb"><ol class="breadcrumb"><li class="breadcrumb-item"><a href="../docs/fsdp_qlora.html">Advanced Features</a></li><li class="breadcrumb-item"><a href="../docs/custom_integrations.html">Custom Integrations</a></li></ol></nav>
<a class="flex-grow-1" role="navigation" data-bs-toggle="collapse" data-bs-target=".quarto-sidebar-collapse-item" aria-controls="quarto-sidebar" aria-expanded="false" aria-label="Toggle sidebar navigation" onclick="if (window.quartoToggleHeadroom) { window.quartoToggleHeadroom(); }">
</a>
</div>
</nav>
</header>
<!-- content -->
<div id="quarto-content" class="quarto-container page-columns page-rows-contents page-layout-article page-navbar">
<!-- sidebar -->
<nav id="quarto-sidebar" class="sidebar collapse collapse-horizontal quarto-sidebar-collapse-item sidebar-navigation docked overflow-auto">
<div class="sidebar-menu-container">
<ul class="list-unstyled mt-1">
<li class="sidebar-item">
<div class="sidebar-item-container">
<a href="../index.html" class="sidebar-item-text sidebar-link">
<span class="menu-text">Home</span></a>
</div>
</li>
<li class="sidebar-item sidebar-item-section">
<div class="sidebar-item-container">
<a class="sidebar-item-text sidebar-link text-start" data-bs-toggle="collapse" data-bs-target="#quarto-sidebar-section-1" role="navigation" aria-expanded="true">
<span class="menu-text">Getting Started</span></a>
<a class="sidebar-item-toggle text-start" data-bs-toggle="collapse" data-bs-target="#quarto-sidebar-section-1" role="navigation" aria-expanded="true" aria-label="Toggle section">
<i class="bi bi-chevron-right ms-2"></i>
</a>
</div>
<ul id="quarto-sidebar-section-1" class="collapse list-unstyled sidebar-section depth1 show">
<li class="sidebar-item">
<div class="sidebar-item-container">
<a href="../docs/getting-started.html" class="sidebar-item-text sidebar-link">
<span class="menu-text">Quickstart</span></a>
</div>
</li>
<li class="sidebar-item">
<div class="sidebar-item-container">
<a href="../docs/choosing_method.html" class="sidebar-item-text sidebar-link">
<span class="menu-text">Which Fine-Tuning Method Should I Use?</span></a>
</div>
</li>
<li class="sidebar-item">
<div class="sidebar-item-container">
<a href="../docs/installation.html" class="sidebar-item-text sidebar-link">
<span class="menu-text">Installation</span></a>
</div>
</li>
<li class="sidebar-item">
<div class="sidebar-item-container">
<a href="../docs/inference.html" class="sidebar-item-text sidebar-link">
<span class="menu-text">Inference and Merging</span></a>
</div>
</li>
<li class="sidebar-item sidebar-item-section">
<div class="sidebar-item-container">
<a class="sidebar-item-text sidebar-link text-start collapsed" data-bs-toggle="collapse" data-bs-target="#quarto-sidebar-section-2" role="navigation" aria-expanded="false">
<span class="menu-text">Model Guides</span></a>
<a class="sidebar-item-toggle text-start collapsed" data-bs-toggle="collapse" data-bs-target="#quarto-sidebar-section-2" role="navigation" aria-expanded="false" aria-label="Toggle section">
<i class="bi bi-chevron-right ms-2"></i>
</a>
</div>
<ul id="quarto-sidebar-section-2" class="collapse list-unstyled sidebar-section depth2 ">
<li class="sidebar-item">
<div class="sidebar-item-container">
<a href="../docs/models/kimi-linear.html" class="sidebar-item-text sidebar-link">
<span class="menu-text">Kimi Linear</span></a>
</div>
</li>
<li class="sidebar-item">
<div class="sidebar-item-container">
<a href="../docs/models/plano.html" class="sidebar-item-text sidebar-link">
<span class="menu-text">Plano Orchestrator</span></a>
</div>
</li>
<li class="sidebar-item">
<div class="sidebar-item-container">
<a href="../docs/models/mimo.html" class="sidebar-item-text sidebar-link">
<span class="menu-text">MiMo</span></a>
</div>
</li>
<li class="sidebar-item">
<div class="sidebar-item-container">
<a href="../docs/models/internvl3_5.html" class="sidebar-item-text sidebar-link">
<span class="menu-text">InternVL 3.5</span></a>
</div>
</li>
<li class="sidebar-item">
<div class="sidebar-item-container">
<a href="../docs/models/olmo3.html" class="sidebar-item-text sidebar-link">
<span class="menu-text">OLMo 3</span></a>
</div>
</li>
<li class="sidebar-item">
<div class="sidebar-item-container">
<a href="../docs/models/trinity.html" class="sidebar-item-text sidebar-link">
<span class="menu-text">Trinity</span></a>
</div>
</li>
<li class="sidebar-item">
<div class="sidebar-item-container">
<a href="../docs/models/arcee.html" class="sidebar-item-text sidebar-link">
<span class="menu-text">Arcee AFM</span></a>
</div>
</li>
<li class="sidebar-item sidebar-item-section">
<div class="sidebar-item-container">
<a class="sidebar-item-text sidebar-link text-start collapsed" data-bs-toggle="collapse" data-bs-target="#quarto-sidebar-section-3" role="navigation" aria-expanded="false">
<span class="menu-text">Ministral3</span></a>
<a class="sidebar-item-toggle text-start collapsed" data-bs-toggle="collapse" data-bs-target="#quarto-sidebar-section-3" role="navigation" aria-expanded="false" aria-label="Toggle section">
<i class="bi bi-chevron-right ms-2"></i>
</a>
</div>
<ul id="quarto-sidebar-section-3" class="collapse list-unstyled sidebar-section depth3 ">
<li class="sidebar-item">
<div class="sidebar-item-container">
<a href="../docs/models/ministral3.html" class="sidebar-item-text sidebar-link">
<span class="menu-text">Ministral3</span></a>
</div>
</li>
<li class="sidebar-item">
<div class="sidebar-item-container">
<a href="../docs/models/ministral3/think.html" class="sidebar-item-text sidebar-link">
<span class="menu-text">Ministral 3 Thinking</span></a>
</div>
</li>
<li class="sidebar-item">
<div class="sidebar-item-container">
<a href="../docs/models/ministral3/vision.html" class="sidebar-item-text sidebar-link">
<span class="menu-text">Ministral 3 Vision</span></a>
</div>
</li>
</ul>
</li>
<li class="sidebar-item sidebar-item-section">
<div class="sidebar-item-container">
<a class="sidebar-item-text sidebar-link text-start collapsed" data-bs-toggle="collapse" data-bs-target="#quarto-sidebar-section-4" role="navigation" aria-expanded="false">
<span class="menu-text">Magistral</span></a>
<a class="sidebar-item-toggle text-start collapsed" data-bs-toggle="collapse" data-bs-target="#quarto-sidebar-section-4" role="navigation" aria-expanded="false" aria-label="Toggle section">
<i class="bi bi-chevron-right ms-2"></i>
</a>
</div>
<ul id="quarto-sidebar-section-4" class="collapse list-unstyled sidebar-section depth3 ">
<li class="sidebar-item">
<div class="sidebar-item-container">
<a href="../docs/models/magistral.html" class="sidebar-item-text sidebar-link">
<span class="menu-text">Magistral</span></a>
</div>
</li>
<li class="sidebar-item">
<div class="sidebar-item-container">
<a href="../docs/models/magistral/think.html" class="sidebar-item-text sidebar-link">
<span class="menu-text">Magistral Thinking</span></a>
</div>
</li>
<li class="sidebar-item">
<div class="sidebar-item-container">
<a href="../docs/models/magistral/vision.html" class="sidebar-item-text sidebar-link">
<span class="menu-text">Magistral Vision</span></a>
</div>
</li>
</ul>
</li>
<li class="sidebar-item">
<div class="sidebar-item-container">
<a href="../docs/models/ministral.html" class="sidebar-item-text sidebar-link">
<span class="menu-text">Ministral</span></a>
</div>
</li>
<li class="sidebar-item">
<div class="sidebar-item-container">
<a href="../docs/models/mistral-small.html" class="sidebar-item-text sidebar-link">
<span class="menu-text">Mistral Small 3.1/3.2</span></a>
</div>
</li>
<li class="sidebar-item">
<div class="sidebar-item-container">
<a href="../docs/models/voxtral.html" class="sidebar-item-text sidebar-link">
<span class="menu-text">Voxtral</span></a>
</div>
</li>
<li class="sidebar-item">
<div class="sidebar-item-container">
<a href="../docs/models/devstral.html" class="sidebar-item-text sidebar-link">
<span class="menu-text">Devstral</span></a>
</div>
</li>
<li class="sidebar-item">
<div class="sidebar-item-container">
<a href="../docs/models/mistral.html" class="sidebar-item-text sidebar-link">
<span class="menu-text">Mistral 7B</span></a>
</div>
</li>
<li class="sidebar-item">
<div class="sidebar-item-container">
<a href="../docs/models/llama-4.html" class="sidebar-item-text sidebar-link">
<span class="menu-text">Llama 4</span></a>
</div>
</li>
<li class="sidebar-item">
<div class="sidebar-item-container">
<a href="../docs/models/llama-2.html" class="sidebar-item-text sidebar-link">
<span class="menu-text">Llama 2</span></a>
</div>
</li>
<li class="sidebar-item">
<div class="sidebar-item-container">
<a href="../docs/models/qwen3-next.html" class="sidebar-item-text sidebar-link">
<span class="menu-text">Qwen 3 Next</span></a>
</div>
</li>
<li class="sidebar-item">
<div class="sidebar-item-container">
<a href="../docs/models/qwen3.html" class="sidebar-item-text sidebar-link">
<span class="menu-text">Qwen 3</span></a>
</div>
</li>
<li class="sidebar-item">
<div class="sidebar-item-container">
<a href="../docs/models/gemma3n.html" class="sidebar-item-text sidebar-link">
<span class="menu-text">Gemma 3n</span></a>
</div>
</li>
<li class="sidebar-item">
<div class="sidebar-item-container">
<a href="../docs/models/apertus.html" class="sidebar-item-text sidebar-link">
<span class="menu-text">Apertus</span></a>
</div>
</li>
<li class="sidebar-item">
<div class="sidebar-item-container">
<a href="../docs/models/gpt-oss.html" class="sidebar-item-text sidebar-link">
<span class="menu-text">GPT-OSS</span></a>
</div>
</li>
<li class="sidebar-item">
<div class="sidebar-item-container">
<a href="../docs/models/seed-oss.html" class="sidebar-item-text sidebar-link">
<span class="menu-text">Seed-OSS</span></a>
</div>
</li>
<li class="sidebar-item">
<div class="sidebar-item-container">
<a href="../docs/models/phi.html" class="sidebar-item-text sidebar-link">
<span class="menu-text">Phi</span></a>
</div>
</li>
<li class="sidebar-item">
<div class="sidebar-item-container">
<a href="../docs/models/smolvlm2.html" class="sidebar-item-text sidebar-link">
<span class="menu-text">SmolVLM 2</span></a>
</div>
</li>
<li class="sidebar-item">
<div class="sidebar-item-container">
<a href="../docs/models/granite4.html" class="sidebar-item-text sidebar-link">
<span class="menu-text">Granite 4</span></a>
</div>
</li>
<li class="sidebar-item">
<div class="sidebar-item-container">
<a href="../docs/models/LiquidAI.html" class="sidebar-item-text sidebar-link">
<span class="menu-text">Liquid Foundation Models 2</span></a>
</div>
</li>
<li class="sidebar-item">
<div class="sidebar-item-container">
<a href="../docs/models/hunyuan.html" class="sidebar-item-text sidebar-link">
<span class="menu-text">Hunyuan</span></a>
</div>
</li>
<li class="sidebar-item">
<div class="sidebar-item-container">
<a href="../docs/models/jamba.html" class="sidebar-item-text sidebar-link">
<span class="menu-text">Jamba</span></a>
</div>
</li>
<li class="sidebar-item">
<div class="sidebar-item-container">
<a href="../docs/models/orpheus.html" class="sidebar-item-text sidebar-link">
<span class="menu-text">Orpheus</span></a>
</div>
</li>
</ul>
</li>
<li class="sidebar-item">
<div class="sidebar-item-container">
<a href="../docs/cli.html" class="sidebar-item-text sidebar-link">
<span class="menu-text">Command Line Interface (CLI)</span></a>
</div>
</li>
<li class="sidebar-item">
<div class="sidebar-item-container">
<a href="../docs/telemetry.html" class="sidebar-item-text sidebar-link">
<span class="menu-text">Telemetry</span></a>
</div>
</li>
<li class="sidebar-item">
<div class="sidebar-item-container">
<a href="../docs/config-reference.html" class="sidebar-item-text sidebar-link">
<span class="menu-text">Config Reference</span></a>
</div>
</li>
<li class="sidebar-item">
<div class="sidebar-item-container">
<a href="../docs/api" class="sidebar-item-text sidebar-link">
<span class="menu-text">API Reference</span></a>
</div>
</li>
</ul>
</li>
<li class="sidebar-item sidebar-item-section">
<div class="sidebar-item-container">
<a href="../docs/dataset-formats/index.html" class="sidebar-item-text sidebar-link">
<span class="menu-text">Dataset Formats</span></a>
<a class="sidebar-item-toggle text-start" data-bs-toggle="collapse" data-bs-target="#quarto-sidebar-section-5" role="navigation" aria-expanded="true" aria-label="Toggle section">
<i class="bi bi-chevron-right ms-2"></i>
</a>
</div>
<ul id="quarto-sidebar-section-5" class="collapse list-unstyled sidebar-section depth1 show">
<li class="sidebar-item">
<div class="sidebar-item-container">
<a href="../docs/dataset-formats/pretraining.html" class="sidebar-item-text sidebar-link">
<span class="menu-text">Pre-training</span></a>
</div>
</li>
<li class="sidebar-item">
<div class="sidebar-item-container">
<a href="../docs/dataset-formats/inst_tune.html" class="sidebar-item-text sidebar-link">
<span class="menu-text">Instruction Tuning</span></a>
</div>
</li>
<li class="sidebar-item">
<div class="sidebar-item-container">
<a href="../docs/dataset-formats/conversation.html" class="sidebar-item-text sidebar-link">
<span class="menu-text">Conversation</span></a>
</div>
</li>
<li class="sidebar-item">
<div class="sidebar-item-container">
<a href="../docs/dataset-formats/stepwise_supervised.html" class="sidebar-item-text sidebar-link">
<span class="menu-text">Stepwise Supervised Format</span></a>
</div>
</li>
<li class="sidebar-item">
<div class="sidebar-item-container">
<a href="../docs/dataset-formats/template_free.html" class="sidebar-item-text sidebar-link">
<span class="menu-text">Template-Free</span></a>
</div>
</li>
<li class="sidebar-item">
<div class="sidebar-item-container">
<a href="../docs/dataset-formats/tokenized.html" class="sidebar-item-text sidebar-link">
<span class="menu-text">Custom Pre-Tokenized Dataset</span></a>
</div>
</li>
</ul>
</li>
<li class="sidebar-item sidebar-item-section">
<div class="sidebar-item-container">
<a class="sidebar-item-text sidebar-link text-start" data-bs-toggle="collapse" data-bs-target="#quarto-sidebar-section-6" role="navigation" aria-expanded="true">
<span class="menu-text">Deployments</span></a>
<a class="sidebar-item-toggle text-start" data-bs-toggle="collapse" data-bs-target="#quarto-sidebar-section-6" role="navigation" aria-expanded="true" aria-label="Toggle section">
<i class="bi bi-chevron-right ms-2"></i>
</a>
</div>
<ul id="quarto-sidebar-section-6" class="collapse list-unstyled sidebar-section depth1 show">
<li class="sidebar-item">
<div class="sidebar-item-container">
<a href="../docs/docker.html" class="sidebar-item-text sidebar-link">
<span class="menu-text">Docker</span></a>
</div>
</li>
<li class="sidebar-item">
<div class="sidebar-item-container">
<a href="../docs/multi-gpu.html" class="sidebar-item-text sidebar-link">
<span class="menu-text">Multi-GPU</span></a>
</div>
</li>
<li class="sidebar-item">
<div class="sidebar-item-container">
<a href="../docs/multi-node.html" class="sidebar-item-text sidebar-link">
<span class="menu-text">Multi Node</span></a>
</div>
</li>
<li class="sidebar-item">
<div class="sidebar-item-container">
<a href="../docs/ray-integration.html" class="sidebar-item-text sidebar-link">
<span class="menu-text">Ray Train</span></a>
</div>
</li>
<li class="sidebar-item">
<div class="sidebar-item-container">
<a href="../docs/amd_hpc.html" class="sidebar-item-text sidebar-link">
<span class="menu-text">AMD GPUs on HPC Systems</span></a>
</div>
</li>
<li class="sidebar-item">
<div class="sidebar-item-container">
<a href="../docs/mac.html" class="sidebar-item-text sidebar-link">
<span class="menu-text">Mac M-series</span></a>
</div>
</li>
</ul>
</li>
<li class="sidebar-item sidebar-item-section">
<div class="sidebar-item-container">
<a class="sidebar-item-text sidebar-link text-start" data-bs-toggle="collapse" data-bs-target="#quarto-sidebar-section-7" role="navigation" aria-expanded="true">
<span class="menu-text">How To Guides</span></a>
<a class="sidebar-item-toggle text-start" data-bs-toggle="collapse" data-bs-target="#quarto-sidebar-section-7" role="navigation" aria-expanded="true" aria-label="Toggle section">
<i class="bi bi-chevron-right ms-2"></i>
</a>
</div>
<ul id="quarto-sidebar-section-7" class="collapse list-unstyled sidebar-section depth1 show">
<li class="sidebar-item">
<div class="sidebar-item-container">
<a href="../docs/multimodal.html" class="sidebar-item-text sidebar-link">
<span class="menu-text">MultiModal / Vision Language Models (BETA)</span></a>
</div>
</li>
<li class="sidebar-item">
<div class="sidebar-item-container">
<a href="../docs/rlhf.html" class="sidebar-item-text sidebar-link">
<span class="menu-text">RLHF (Beta)</span></a>
</div>
</li>
<li class="sidebar-item">
<div class="sidebar-item-container">
<a href="../docs/grpo.html" class="sidebar-item-text sidebar-link">
<span class="menu-text">GRPO Training</span></a>
</div>
</li>
<li class="sidebar-item">
<div class="sidebar-item-container">
<a href="../docs/ebft.html" class="sidebar-item-text sidebar-link">
<span class="menu-text">EBFT Training</span></a>
</div>
</li>
<li class="sidebar-item">
<div class="sidebar-item-container">
<a href="../docs/vllm_serving.html" class="sidebar-item-text sidebar-link">
<span class="menu-text">vLLM Serving for GRPO Training</span></a>
</div>
</li>
<li class="sidebar-item">
<div class="sidebar-item-container">
<a href="../docs/reward_modelling.html" class="sidebar-item-text sidebar-link">
<span class="menu-text">Reward Modelling</span></a>
</div>
</li>
<li class="sidebar-item">
<div class="sidebar-item-container">
<a href="../docs/lr_groups.html" class="sidebar-item-text sidebar-link">
<span class="menu-text">Learning Rate Groups</span></a>
</div>
</li>
<li class="sidebar-item">
<div class="sidebar-item-container">
<a href="../docs/lora_optims.html" class="sidebar-item-text sidebar-link">
<span class="menu-text">LoRA Optimizations</span></a>
</div>
</li>
<li class="sidebar-item">
<div class="sidebar-item-container">
<a href="../docs/dataset_loading.html" class="sidebar-item-text sidebar-link">
<span class="menu-text">Dataset Loading</span></a>
</div>
</li>
<li class="sidebar-item">
<div class="sidebar-item-container">
<a href="../docs/qat.html" class="sidebar-item-text sidebar-link">
<span class="menu-text">Quantization Aware Training (QAT)</span></a>
</div>
</li>
<li class="sidebar-item">
<div class="sidebar-item-container">
<a href="../docs/quantize.html" class="sidebar-item-text sidebar-link">
<span class="menu-text">Quantization with torchao</span></a>
</div>
</li>
<li class="sidebar-item">
<div class="sidebar-item-container">
<a href="../docs/optimizations.html" class="sidebar-item-text sidebar-link">
<span class="menu-text">Optimizations Guide</span></a>
</div>
</li>
</ul>
</li>
<li class="sidebar-item sidebar-item-section">
<div class="sidebar-item-container">
<a class="sidebar-item-text sidebar-link text-start" data-bs-toggle="collapse" data-bs-target="#quarto-sidebar-section-8" role="navigation" aria-expanded="true">
<span class="menu-text">Core Concepts</span></a>
<a class="sidebar-item-toggle text-start" data-bs-toggle="collapse" data-bs-target="#quarto-sidebar-section-8" role="navigation" aria-expanded="true" aria-label="Toggle section">
<i class="bi bi-chevron-right ms-2"></i>
</a>
</div>
<ul id="quarto-sidebar-section-8" class="collapse list-unstyled sidebar-section depth1 show">
<li class="sidebar-item">
<div class="sidebar-item-container">
<a href="../docs/batch_vs_grad.html" class="sidebar-item-text sidebar-link">
<span class="menu-text">Batch size vs Gradient accumulation</span></a>
</div>
</li>
<li class="sidebar-item">
<div class="sidebar-item-container">
<a href="../docs/dataset_preprocessing.html" class="sidebar-item-text sidebar-link">
<span class="menu-text">Dataset Preprocessing</span></a>
</div>
</li>
<li class="sidebar-item">
<div class="sidebar-item-container">
<a href="../docs/streaming.html" class="sidebar-item-text sidebar-link">
<span class="menu-text">Streaming Datasets</span></a>
</div>
</li>
<li class="sidebar-item">
<div class="sidebar-item-container">
<a href="../docs/multipack.html" class="sidebar-item-text sidebar-link">
<span class="menu-text">Multipack (Sample Packing)</span></a>
</div>
</li>
<li class="sidebar-item">
<div class="sidebar-item-container">
<a href="../docs/mixed_precision.html" class="sidebar-item-text sidebar-link">
<span class="menu-text">Mixed Precision Training</span></a>
</div>
</li>
<li class="sidebar-item">
<div class="sidebar-item-container">
<a href="../docs/optimizers.html" class="sidebar-item-text sidebar-link">
<span class="menu-text">Optimizers</span></a>
</div>
</li>
<li class="sidebar-item">
<div class="sidebar-item-container">
<a href="../docs/attention.html" class="sidebar-item-text sidebar-link">
<span class="menu-text">Attention</span></a>
</div>
</li>
</ul>
</li>
<li class="sidebar-item sidebar-item-section">
<div class="sidebar-item-container">
<a class="sidebar-item-text sidebar-link text-start" data-bs-toggle="collapse" data-bs-target="#quarto-sidebar-section-9" role="navigation" aria-expanded="true">
<span class="menu-text">Advanced Features</span></a>
<a class="sidebar-item-toggle text-start" data-bs-toggle="collapse" data-bs-target="#quarto-sidebar-section-9" role="navigation" aria-expanded="true" aria-label="Toggle section">
<i class="bi bi-chevron-right ms-2"></i>
</a>
</div>
<ul id="quarto-sidebar-section-9" class="collapse list-unstyled sidebar-section depth1 show">
<li class="sidebar-item">
<div class="sidebar-item-container">
<a href="../docs/fsdp_qlora.html" class="sidebar-item-text sidebar-link">
<span class="menu-text">FSDP + QLoRA</span></a>
</div>
</li>
<li class="sidebar-item">
<div class="sidebar-item-container">
<a href="../docs/torchao.html" class="sidebar-item-text sidebar-link">
<span class="menu-text">PyTorch ao</span></a>
</div>
</li>
<li class="sidebar-item">
<div class="sidebar-item-container">
<a href="../docs/custom_integrations.html" class="sidebar-item-text sidebar-link active">
<span class="menu-text">Custom Integrations</span></a>
</div>
</li>
<li class="sidebar-item">
<div class="sidebar-item-container">
<a href="../docs/sequence_parallelism.html" class="sidebar-item-text sidebar-link">
<span class="menu-text">Sequence Parallelism</span></a>
</div>
</li>
<li class="sidebar-item">
<div class="sidebar-item-container">
<a href="../docs/gradient_checkpointing.html" class="sidebar-item-text sidebar-link">
<span class="menu-text">Gradient Checkpointing, Activation Offloading, and Layer Offloading</span></a>
</div>
</li>
<li class="sidebar-item">
<div class="sidebar-item-container">
<a href="../docs/nd_parallelism.html" class="sidebar-item-text sidebar-link">
<span class="menu-text">N-D Parallelism (Beta)</span></a>
</div>
</li>
<li class="sidebar-item">
<div class="sidebar-item-container">
<a href="../docs/expert_quantization.html" class="sidebar-item-text sidebar-link">
<span class="menu-text">MoE Expert Quantization</span></a>
</div>
</li>
</ul>
</li>
<li class="sidebar-item sidebar-item-section">
<div class="sidebar-item-container">
<a class="sidebar-item-text sidebar-link text-start" data-bs-toggle="collapse" data-bs-target="#quarto-sidebar-section-10" role="navigation" aria-expanded="true">
<span class="menu-text">Troubleshooting</span></a>
<a class="sidebar-item-toggle text-start" data-bs-toggle="collapse" data-bs-target="#quarto-sidebar-section-10" role="navigation" aria-expanded="true" aria-label="Toggle section">
<i class="bi bi-chevron-right ms-2"></i>
</a>
</div>
<ul id="quarto-sidebar-section-10" class="collapse list-unstyled sidebar-section depth1 show">
<li class="sidebar-item">
<div class="sidebar-item-container">
<a href="../docs/faq.html" class="sidebar-item-text sidebar-link">
<span class="menu-text">FAQ</span></a>
</div>
</li>
<li class="sidebar-item">
<div class="sidebar-item-container">
<a href="../docs/training_stability.html" class="sidebar-item-text sidebar-link">
<span class="menu-text">Training Stability &amp; Debugging</span></a>
</div>
</li>
<li class="sidebar-item">
<div class="sidebar-item-container">
<a href="../docs/debugging.html" class="sidebar-item-text sidebar-link">
<span class="menu-text">Debugging</span></a>
</div>
</li>
<li class="sidebar-item">
<div class="sidebar-item-container">
<a href="../docs/nccl.html" class="sidebar-item-text sidebar-link">
<span class="menu-text">NCCL</span></a>
</div>
</li>
</ul>
</li>
</ul>
</div>
</nav>
<div id="quarto-sidebar-glass" class="quarto-sidebar-collapse-item" data-bs-toggle="collapse" data-bs-target=".quarto-sidebar-collapse-item"></div>
<!-- margin-sidebar -->
<div id="quarto-margin-sidebar" class="sidebar margin-sidebar">
<nav id="TOC" role="doc-toc" class="toc-active">
<h2 id="toc-title">On this page</h2>
<ul>
<li><a href="#cut-cross-entropy" id="toc-cut-cross-entropy" class="nav-link active" data-scroll-target="#cut-cross-entropy">Cut Cross Entropy</a>
<ul class="collapse">
<li><a href="#requirements" id="toc-requirements" class="nav-link" data-scroll-target="#requirements">Requirements</a></li>
<li><a href="#installation" id="toc-installation" class="nav-link" data-scroll-target="#installation">Installation</a></li>
<li><a href="#usage" id="toc-usage" class="nav-link" data-scroll-target="#usage">Usage</a></li>
<li><a href="#supported-models" id="toc-supported-models" class="nav-link" data-scroll-target="#supported-models">Supported Models</a></li>
<li><a href="#citation" id="toc-citation" class="nav-link" data-scroll-target="#citation">Citation</a></li>
</ul></li>
<li><a href="#densemixer" id="toc-densemixer" class="nav-link" data-scroll-target="#densemixer">DenseMixer</a></li>
<li><a href="#diffusion-lm-training-plugin-for-axolotl" id="toc-diffusion-lm-training-plugin-for-axolotl" class="nav-link" data-scroll-target="#diffusion-lm-training-plugin-for-axolotl">Diffusion LM Training Plugin for Axolotl</a>
<ul class="collapse">
<li><a href="#overview" id="toc-overview" class="nav-link" data-scroll-target="#overview">Overview</a></li>
<li><a href="#installation-1" id="toc-installation-1" class="nav-link" data-scroll-target="#installation-1">Installation</a></li>
<li><a href="#quickstart" id="toc-quickstart" class="nav-link" data-scroll-target="#quickstart">Quickstart</a></li>
<li><a href="#basic-configuration" id="toc-basic-configuration" class="nav-link" data-scroll-target="#basic-configuration">Basic Configuration</a></li>
<li><a href="#supported-models-1" id="toc-supported-models-1" class="nav-link" data-scroll-target="#supported-models-1">Supported Models</a></li>
<li><a href="#how-it-works" id="toc-how-it-works" class="nav-link" data-scroll-target="#how-it-works">How It Works</a></li>
<li><a href="#random-masking" id="toc-random-masking" class="nav-link" data-scroll-target="#random-masking">Random Masking</a></li>
<li><a href="#diffusion-loss" id="toc-diffusion-loss" class="nav-link" data-scroll-target="#diffusion-loss">Diffusion Loss</a></li>
<li><a href="#sample-generation" id="toc-sample-generation" class="nav-link" data-scroll-target="#sample-generation">Sample Generation</a></li>
<li><a href="#inference" id="toc-inference" class="nav-link" data-scroll-target="#inference">Inference</a></li>
<li><a href="#metrics-and-monitoring" id="toc-metrics-and-monitoring" class="nav-link" data-scroll-target="#metrics-and-monitoring">Metrics and Monitoring</a></li>
<li><a href="#limitations" id="toc-limitations" class="nav-link" data-scroll-target="#limitations">Limitations</a></li>
<li><a href="#references" id="toc-references" class="nav-link" data-scroll-target="#references">References</a></li>
</ul></li>
<li><a href="#grokfast" id="toc-grokfast" class="nav-link" data-scroll-target="#grokfast">Grokfast</a>
<ul class="collapse">
<li><a href="#usage-1" id="toc-usage-1" class="nav-link" data-scroll-target="#usage-1">Usage</a></li>
<li><a href="#citation-1" id="toc-citation-1" class="nav-link" data-scroll-target="#citation-1">Citation</a></li>
</ul></li>
<li><a href="#kernels-integration" id="toc-kernels-integration" class="nav-link" data-scroll-target="#kernels-integration">Kernels Integration</a>
<ul class="collapse">
<li><a href="#usage-2" id="toc-usage-2" class="nav-link" data-scroll-target="#usage-2">Usage</a></li>
<li><a href="#sonicmoe-installation" id="toc-sonicmoe-installation" class="nav-link" data-scroll-target="#sonicmoe-installation">SonicMoE installation</a></li>
<li><a href="#how-it-works-1" id="toc-how-it-works-1" class="nav-link" data-scroll-target="#how-it-works-1">How It Works</a></li>
<li><a href="#scattermoe" id="toc-scattermoe" class="nav-link" data-scroll-target="#scattermoe">ScatterMoE</a></li>
<li><a href="#sonicmoe" id="toc-sonicmoe" class="nav-link" data-scroll-target="#sonicmoe">SonicMoE</a></li>
<li><a href="#model-support-matrix" id="toc-model-support-matrix" class="nav-link" data-scroll-target="#model-support-matrix">Model Support Matrix</a></li>
<li><a href="#routing-strategies" id="toc-routing-strategies" class="nav-link" data-scroll-target="#routing-strategies">Routing strategies</a></li>
<li><a href="#per-model-support" id="toc-per-model-support" class="nav-link" data-scroll-target="#per-model-support">Per-model support</a></li>
<li><a href="#feature-comparison" id="toc-feature-comparison" class="nav-link" data-scroll-target="#feature-comparison">Feature comparison</a></li>
<li><a href="#shared-expert-handling" id="toc-shared-expert-handling" class="nav-link" data-scroll-target="#shared-expert-handling">Shared Expert Handling</a></li>
<li><a href="#gemma-4" id="toc-gemma-4" class="nav-link" data-scroll-target="#gemma-4">Gemma 4</a></li>
<li><a href="#limitations-1" id="toc-limitations-1" class="nav-link" data-scroll-target="#limitations-1">Limitations</a></li>
<li><a href="#note-on-megablocks" id="toc-note-on-megablocks" class="nav-link" data-scroll-target="#note-on-megablocks">Note on MegaBlocks</a></li>
</ul></li>
<li><a href="#knowledge-distillation-kd" id="toc-knowledge-distillation-kd" class="nav-link" data-scroll-target="#knowledge-distillation-kd">Knowledge Distillation (KD)</a>
<ul class="collapse">
<li><a href="#usage-3" id="toc-usage-3" class="nav-link" data-scroll-target="#usage-3">Usage</a></li>
</ul></li>
<li><a href="#llmcompressor" id="toc-llmcompressor" class="nav-link" data-scroll-target="#llmcompressor">LLMCompressor</a>
<ul class="collapse">
<li><a href="#requirements-1" id="toc-requirements-1" class="nav-link" data-scroll-target="#requirements-1">Requirements</a></li>
<li><a href="#usage-4" id="toc-usage-4" class="nav-link" data-scroll-target="#usage-4">Usage</a></li>
<li><a href="#storage-optimization-with-save_compressed" id="toc-storage-optimization-with-save_compressed" class="nav-link" data-scroll-target="#storage-optimization-with-save_compressed">Storage Optimization with save_compressed</a></li>
<li><a href="#example-config" id="toc-example-config" class="nav-link" data-scroll-target="#example-config">Example Config</a></li>
<li><a href="#inference-with-vllm" id="toc-inference-with-vllm" class="nav-link" data-scroll-target="#inference-with-vllm">Inference with vLLM</a></li>
<li><a href="#learn-more" id="toc-learn-more" class="nav-link" data-scroll-target="#learn-more">Learn More</a></li>
</ul></li>
<li><a href="#language-model-evaluation-harness-lm-eval" id="toc-language-model-evaluation-harness-lm-eval" class="nav-link" data-scroll-target="#language-model-evaluation-harness-lm-eval">Language Model Evaluation Harness (LM Eval)</a>
<ul class="collapse">
<li><a href="#usage-5" id="toc-usage-5" class="nav-link" data-scroll-target="#usage-5">Usage</a></li>
<li><a href="#post-training-evaluation" id="toc-post-training-evaluation" class="nav-link" data-scroll-target="#post-training-evaluation">1. Post-Training Evaluation</a></li>
<li><a href="#standalone-cli-evaluation" id="toc-standalone-cli-evaluation" class="nav-link" data-scroll-target="#standalone-cli-evaluation">2. Standalone CLI Evaluation</a></li>
<li><a href="#model-selection-priority" id="toc-model-selection-priority" class="nav-link" data-scroll-target="#model-selection-priority">Model Selection Priority</a></li>
<li><a href="#citation-2" id="toc-citation-2" class="nav-link" data-scroll-target="#citation-2">Citation</a></li>
</ul></li>
<li><a href="#liger-kernels" id="toc-liger-kernels" class="nav-link" data-scroll-target="#liger-kernels">Liger Kernels</a>
<ul class="collapse">
<li><a href="#usage-6" id="toc-usage-6" class="nav-link" data-scroll-target="#usage-6">Usage</a></li>
<li><a href="#supported-models-2" id="toc-supported-models-2" class="nav-link" data-scroll-target="#supported-models-2">Supported Models</a></li>
<li><a href="#citation-3" id="toc-citation-3" class="nav-link" data-scroll-target="#citation-3">Citation</a></li>
</ul></li>
<li><a href="#nemo-gym-integration-for-axolotl" id="toc-nemo-gym-integration-for-axolotl" class="nav-link" data-scroll-target="#nemo-gym-integration-for-axolotl">NeMo Gym Integration for Axolotl</a>
<ul class="collapse">
<li><a href="#validated-training-paths" id="toc-validated-training-paths" class="nav-link" data-scroll-target="#validated-training-paths">Validated Training Paths</a></li>
<li><a href="#quick-start" id="toc-quick-start" class="nav-link" data-scroll-target="#quick-start">Quick Start</a></li>
<li><a href="#prerequisites" id="toc-prerequisites" class="nav-link" data-scroll-target="#prerequisites">Prerequisites</a></li>
<li><a href="#set-up-nemo-gym" id="toc-set-up-nemo-gym" class="nav-link" data-scroll-target="#set-up-nemo-gym">1. Set Up NeMo Gym</a></li>
<li><a href="#multi-turn-with-async-grpo-recommended-fastest-path" id="toc-multi-turn-with-async-grpo-recommended-fastest-path" class="nav-link" data-scroll-target="#multi-turn-with-async-grpo-recommended-fastest-path">2. Multi-Turn with Async GRPO (Recommended — Fastest Path)</a></li>
<li><a href="#single-turn-training-simplest-no-agent-server-needed" id="toc-single-turn-training-simplest-no-agent-server-needed" class="nav-link" data-scroll-target="#single-turn-training-simplest-no-agent-server-needed">3. Single-Turn Training (Simplest — No Agent Server Needed)</a></li>
<li><a href="#how-it-works-2" id="toc-how-it-works-2" class="nav-link" data-scroll-target="#how-it-works-2">How It Works</a></li>
<li><a href="#single-turn" id="toc-single-turn" class="nav-link" data-scroll-target="#single-turn">Single-Turn</a></li>
<li><a href="#multi-turn-agent-run" id="toc-multi-turn-agent-run" class="nav-link" data-scroll-target="#multi-turn-agent-run">Multi-Turn (Agent /run)</a></li>
<li><a href="#data-producer-architecture-multi-turn" id="toc-data-producer-architecture-multi-turn" class="nav-link" data-scroll-target="#data-producer-architecture-multi-turn">Data Producer Architecture (Multi-Turn)</a></li>
<li><a href="#weight-sync-lora-mode" id="toc-weight-sync-lora-mode" class="nav-link" data-scroll-target="#weight-sync-lora-mode">Weight Sync (LoRA Mode)</a></li>
<li><a href="#multi-environment-support" id="toc-multi-environment-support" class="nav-link" data-scroll-target="#multi-environment-support">Multi-Environment Support</a></li>
<li><a href="#configuration-reference" id="toc-configuration-reference" class="nav-link" data-scroll-target="#configuration-reference">Configuration Reference</a></li>
<li><a href="#dataset-jsonl-format" id="toc-dataset-jsonl-format" class="nav-link" data-scroll-target="#dataset-jsonl-format">Dataset JSONL Format</a></li>
<li><a href="#reward-functions" id="toc-reward-functions" class="nav-link" data-scroll-target="#reward-functions">Reward Functions</a></li>
<li><a href="#known-issues-troubleshooting" id="toc-known-issues-troubleshooting" class="nav-link" data-scroll-target="#known-issues-troubleshooting">Known Issues / Troubleshooting</a></li>
<li><a href="#nemo-gym-server-setup" id="toc-nemo-gym-server-setup" class="nav-link" data-scroll-target="#nemo-gym-server-setup">NeMo Gym Server Setup</a></li>
<li><a href="#vllm-weight-sync" id="toc-vllm-weight-sync" class="nav-link" data-scroll-target="#vllm-weight-sync">vLLM / Weight Sync</a></li>
<li><a href="#multi-turn" id="toc-multi-turn" class="nav-link" data-scroll-target="#multi-turn">Multi-Turn</a></li>
<li><a href="#fsdp2" id="toc-fsdp2" class="nav-link" data-scroll-target="#fsdp2">FSDP2</a></li>
<li><a href="#comparison-with-other-integrations" id="toc-comparison-with-other-integrations" class="nav-link" data-scroll-target="#comparison-with-other-integrations">Comparison with Other Integrations</a></li>
</ul></li>
<li><a href="#spectrum" id="toc-spectrum" class="nav-link" data-scroll-target="#spectrum">Spectrum</a>
<ul class="collapse">
<li><a href="#overview-1" id="toc-overview-1" class="nav-link" data-scroll-target="#overview-1">Overview</a></li>
<li><a href="#usage-7" id="toc-usage-7" class="nav-link" data-scroll-target="#usage-7">Usage</a></li>
<li><a href="#citation-4" id="toc-citation-4" class="nav-link" data-scroll-target="#citation-4">Citation</a></li>
</ul></li>
<li><a href="#swanlab-integration-for-axolotl" id="toc-swanlab-integration-for-axolotl" class="nav-link" data-scroll-target="#swanlab-integration-for-axolotl">SwanLab Integration for Axolotl</a>
<ul class="collapse">
<li><a href="#features" id="toc-features" class="nav-link" data-scroll-target="#features">Features</a></li>
<li><a href="#installation-2" id="toc-installation-2" class="nav-link" data-scroll-target="#installation-2">Installation</a></li>
<li><a href="#quick-start-1" id="toc-quick-start-1" class="nav-link" data-scroll-target="#quick-start-1">Quick Start</a></li>
<li><a href="#register-for-swanlab-optional-for-cloud-mode" id="toc-register-for-swanlab-optional-for-cloud-mode" class="nav-link" data-scroll-target="#register-for-swanlab-optional-for-cloud-mode">1. Register for SwanLab (Optional for cloud mode)</a></li>
<li><a href="#configure-axolotl-config-file" id="toc-configure-axolotl-config-file" class="nav-link" data-scroll-target="#configure-axolotl-config-file">2. Configure Axolotl Config File</a></li>
<li><a href="#run-training" id="toc-run-training" class="nav-link" data-scroll-target="#run-training">3. Run Training</a></li>
<li><a href="#configuration-options" id="toc-configuration-options" class="nav-link" data-scroll-target="#configuration-options">Configuration Options</a></li>
<li><a href="#basic-configuration-1" id="toc-basic-configuration-1" class="nav-link" data-scroll-target="#basic-configuration-1">Basic Configuration</a></li>
<li><a href="#advanced-configuration" id="toc-advanced-configuration" class="nav-link" data-scroll-target="#advanced-configuration">Advanced Configuration</a></li>
<li><a href="#configuration-examples" id="toc-configuration-examples" class="nav-link" data-scroll-target="#configuration-examples">Configuration Examples</a></li>
<li><a href="#example-1-basic-cloud-sync" id="toc-example-1-basic-cloud-sync" class="nav-link" data-scroll-target="#example-1-basic-cloud-sync">Example 1: Basic Cloud Sync</a></li>
<li><a href="#example-2-offlinelocal-mode" id="toc-example-2-offlinelocal-mode" class="nav-link" data-scroll-target="#example-2-offlinelocal-mode">Example 2: Offline/Local Mode</a></li>
<li><a href="#example-3-team-workspace" id="toc-example-3-team-workspace" class="nav-link" data-scroll-target="#example-3-team-workspace">Example 3: Team Workspace</a></li>
<li><a href="#example-4-private-deployment" id="toc-example-4-private-deployment" class="nav-link" data-scroll-target="#example-4-private-deployment">Example 4: Private Deployment</a></li>
<li><a href="#team-notifications-with-lark-feishu" id="toc-team-notifications-with-lark-feishu" class="nav-link" data-scroll-target="#team-notifications-with-lark-feishu">Team Notifications with Lark (Feishu)</a></li>
<li><a href="#prerequisites-1" id="toc-prerequisites-1" class="nav-link" data-scroll-target="#prerequisites-1">Prerequisites</a></li>
<li><a href="#example-5-basic-lark-notifications" id="toc-example-5-basic-lark-notifications" class="nav-link" data-scroll-target="#example-5-basic-lark-notifications">Example 5: Basic Lark Notifications</a></li>
<li><a href="#example-6-lark-notifications-with-hmac-security-recommended" id="toc-example-6-lark-notifications-with-hmac-security-recommended" class="nav-link" data-scroll-target="#example-6-lark-notifications-with-hmac-security-recommended">Example 6: Lark Notifications with HMAC Security (Recommended)</a></li>
<li><a href="#example-7-team-workspace-lark-notifications" id="toc-example-7-team-workspace-lark-notifications" class="nav-link" data-scroll-target="#example-7-team-workspace-lark-notifications">Example 7: Team Workspace + Lark Notifications</a></li>
<li><a href="#what-notifications-are-sent" id="toc-what-notifications-are-sent" class="nav-link" data-scroll-target="#what-notifications-are-sent">What Notifications Are Sent?</a></li>
<li><a href="#lark-configuration-validation" id="toc-lark-configuration-validation" class="nav-link" data-scroll-target="#lark-configuration-validation">Lark Configuration Validation</a></li>
<li><a href="#security-best-practices" id="toc-security-best-practices" class="nav-link" data-scroll-target="#security-best-practices">Security Best Practices</a></li>
<li><a href="#distributed-training" id="toc-distributed-training" class="nav-link" data-scroll-target="#distributed-training">Distributed Training</a></li>
<li><a href="#rlhf-completion-table-logging" id="toc-rlhf-completion-table-logging" class="nav-link" data-scroll-target="#rlhf-completion-table-logging">RLHF Completion Table Logging</a></li>
<li><a href="#features-1" id="toc-features-1" class="nav-link" data-scroll-target="#features-1">Features</a></li>
<li><a href="#configuration" id="toc-configuration" class="nav-link" data-scroll-target="#configuration">Configuration</a></li>
<li><a href="#example-dpo-training-with-completion-logging" id="toc-example-dpo-training-with-completion-logging" class="nav-link" data-scroll-target="#example-dpo-training-with-completion-logging">Example: DPO Training with Completion Logging</a></li>
<li><a href="#example-disable-completion-logging" id="toc-example-disable-completion-logging" class="nav-link" data-scroll-target="#example-disable-completion-logging">Example: Disable Completion Logging</a></li>
<li><a href="#supported-rlhf-trainers" id="toc-supported-rlhf-trainers" class="nav-link" data-scroll-target="#supported-rlhf-trainers">Supported RLHF Trainers</a></li>
<li><a href="#how-it-works-3" id="toc-how-it-works-3" class="nav-link" data-scroll-target="#how-it-works-3">How It Works</a></li>
<li><a href="#viewing-completion-tables" id="toc-viewing-completion-tables" class="nav-link" data-scroll-target="#viewing-completion-tables">Viewing Completion Tables</a></li>
<li><a href="#memory-management" id="toc-memory-management" class="nav-link" data-scroll-target="#memory-management">Memory Management</a></li>
<li><a href="#performance-impact" id="toc-performance-impact" class="nav-link" data-scroll-target="#performance-impact">Performance Impact</a></li>
<li><a href="#troubleshooting" id="toc-troubleshooting" class="nav-link" data-scroll-target="#troubleshooting">Troubleshooting</a></li>
<li><a href="#performance-profiling" id="toc-performance-profiling" class="nav-link" data-scroll-target="#performance-profiling">Performance Profiling</a></li>
<li><a href="#features-2" id="toc-features-2" class="nav-link" data-scroll-target="#features-2">Features</a></li>
<li><a href="#basic-usage-decorator" id="toc-basic-usage-decorator" class="nav-link" data-scroll-target="#basic-usage-decorator">Basic Usage: Decorator</a></li>
<li><a href="#advanced-usage-context-manager" id="toc-advanced-usage-context-manager" class="nav-link" data-scroll-target="#advanced-usage-context-manager">Advanced Usage: Context Manager</a></li>
<li><a href="#advanced-usage-profilingconfig" id="toc-advanced-usage-profilingconfig" class="nav-link" data-scroll-target="#advanced-usage-profilingconfig">Advanced Usage: ProfilingConfig</a></li>
<li><a href="#viewing-profiling-metrics" id="toc-viewing-profiling-metrics" class="nav-link" data-scroll-target="#viewing-profiling-metrics">Viewing Profiling Metrics</a></li>
<li><a href="#configuration-in-axolotl-config" id="toc-configuration-in-axolotl-config" class="nav-link" data-scroll-target="#configuration-in-axolotl-config">Configuration in Axolotl Config</a></li>
<li><a href="#performance-impact-1" id="toc-performance-impact-1" class="nav-link" data-scroll-target="#performance-impact-1">Performance Impact</a></li>
<li><a href="#best-practices" id="toc-best-practices" class="nav-link" data-scroll-target="#best-practices">Best Practices</a></li>
<li><a href="#example-complete-profiling-setup" id="toc-example-complete-profiling-setup" class="nav-link" data-scroll-target="#example-complete-profiling-setup">Example: Complete Profiling Setup</a></li>
<li><a href="#troubleshooting-1" id="toc-troubleshooting-1" class="nav-link" data-scroll-target="#troubleshooting-1">Troubleshooting</a></li>
<li><a href="#complete-config-example" id="toc-complete-config-example" class="nav-link" data-scroll-target="#complete-config-example">Complete Config Example</a></li>
<li><a href="#modes-explained" id="toc-modes-explained" class="nav-link" data-scroll-target="#modes-explained">Modes Explained</a></li>
<li><a href="#cloud-mode-default" id="toc-cloud-mode-default" class="nav-link" data-scroll-target="#cloud-mode-default"><code>cloud</code> Mode (Default)</a></li>
<li><a href="#local-mode" id="toc-local-mode" class="nav-link" data-scroll-target="#local-mode"><code>local</code> Mode</a></li>
<li><a href="#offline-mode" id="toc-offline-mode" class="nav-link" data-scroll-target="#offline-mode"><code>offline</code> Mode</a></li>
<li><a href="#disabled-mode" id="toc-disabled-mode" class="nav-link" data-scroll-target="#disabled-mode"><code>disabled</code> Mode</a></li>
<li><a href="#configuration-validation-conflict-detection" id="toc-configuration-validation-conflict-detection" class="nav-link" data-scroll-target="#configuration-validation-conflict-detection">Configuration Validation &amp; Conflict Detection</a></li>
<li><a href="#required-fields-validation" id="toc-required-fields-validation" class="nav-link" data-scroll-target="#required-fields-validation">Required Fields Validation</a></li>
<li><a href="#cloud-mode-api-key-warning" id="toc-cloud-mode-api-key-warning" class="nav-link" data-scroll-target="#cloud-mode-api-key-warning">Cloud Mode API Key Warning</a></li>
<li><a href="#multi-logger-performance-warnings" id="toc-multi-logger-performance-warnings" class="nav-link" data-scroll-target="#multi-logger-performance-warnings">Multi-Logger Performance Warnings</a></li>
<li><a href="#auto-enable-logic" id="toc-auto-enable-logic" class="nav-link" data-scroll-target="#auto-enable-logic">Auto-Enable Logic</a></li>
<li><a href="#distributed-training-detection" id="toc-distributed-training-detection" class="nav-link" data-scroll-target="#distributed-training-detection">Distributed Training Detection</a></li>
<li><a href="#authentication" id="toc-authentication" class="nav-link" data-scroll-target="#authentication">Authentication</a></li>
<li><a href="#method-1-environment-variable-recommended" id="toc-method-1-environment-variable-recommended" class="nav-link" data-scroll-target="#method-1-environment-variable-recommended">Method 1: Environment Variable (Recommended)</a></li>
<li><a href="#method-2-login-command" id="toc-method-2-login-command" class="nav-link" data-scroll-target="#method-2-login-command">Method 2: Login Command</a></li>
<li><a href="#method-3-config-file" id="toc-method-3-config-file" class="nav-link" data-scroll-target="#method-3-config-file">Method 3: Config File</a></li>
<li><a href="#what-gets-logged" id="toc-what-gets-logged" class="nav-link" data-scroll-target="#what-gets-logged">What Gets Logged?</a></li>
<li><a href="#automatically-logged-metrics" id="toc-automatically-logged-metrics" class="nav-link" data-scroll-target="#automatically-logged-metrics">Automatically Logged Metrics</a></li>
<li><a href="#automatically-logged-config" id="toc-automatically-logged-config" class="nav-link" data-scroll-target="#automatically-logged-config">Automatically Logged Config</a></li>
<li><a href="#viewing-your-experiments" id="toc-viewing-your-experiments" class="nav-link" data-scroll-target="#viewing-your-experiments">Viewing Your Experiments</a></li>
<li><a href="#cloud-mode" id="toc-cloud-mode" class="nav-link" data-scroll-target="#cloud-mode">Cloud Mode</a></li>
<li><a href="#local-mode-1" id="toc-local-mode-1" class="nav-link" data-scroll-target="#local-mode-1">Local Mode</a></li>
<li><a href="#integration-with-existing-tools" id="toc-integration-with-existing-tools" class="nav-link" data-scroll-target="#integration-with-existing-tools">Integration with Existing Tools</a></li>
<li><a href="#troubleshooting-2" id="toc-troubleshooting-2" class="nav-link" data-scroll-target="#troubleshooting-2">Troubleshooting</a></li>
<li><a href="#configuration-errors" id="toc-configuration-errors" class="nav-link" data-scroll-target="#configuration-errors">Configuration Errors</a></li>
<li><a href="#import-errors" id="toc-import-errors" class="nav-link" data-scroll-target="#import-errors">Import Errors</a></li>
<li><a href="#performance-issues" id="toc-performance-issues" class="nav-link" data-scroll-target="#performance-issues">Performance Issues</a></li>
<li><a href="#distributed-training-issues" id="toc-distributed-training-issues" class="nav-link" data-scroll-target="#distributed-training-issues">Distributed Training Issues</a></li>
<li><a href="#swanlab-not-logging-metrics" id="toc-swanlab-not-logging-metrics" class="nav-link" data-scroll-target="#swanlab-not-logging-metrics">SwanLab not logging metrics</a></li>
<li><a href="#api-key-errors" id="toc-api-key-errors" class="nav-link" data-scroll-target="#api-key-errors">API Key errors</a></li>
<li><a href="#cloud-sync-issues" id="toc-cloud-sync-issues" class="nav-link" data-scroll-target="#cloud-sync-issues">Cloud sync issues</a></li>
<li><a href="#plugin-not-loaded" id="toc-plugin-not-loaded" class="nav-link" data-scroll-target="#plugin-not-loaded">Plugin not loaded</a></li>
<li><a href="#lark-notification-issues" id="toc-lark-notification-issues" class="nav-link" data-scroll-target="#lark-notification-issues">Lark Notification Issues</a></li>
<li><a href="#comparison-swanlab-vs-wandb" id="toc-comparison-swanlab-vs-wandb" class="nav-link" data-scroll-target="#comparison-swanlab-vs-wandb">Comparison: SwanLab vs WandB</a></li>
<li><a href="#advanced-usage" id="toc-advanced-usage" class="nav-link" data-scroll-target="#advanced-usage">Advanced Usage</a></li>
<li><a href="#custom-logging" id="toc-custom-logging" class="nav-link" data-scroll-target="#custom-logging">Custom Logging</a></li>
<li><a href="#experiment-comparison" id="toc-experiment-comparison" class="nav-link" data-scroll-target="#experiment-comparison">Experiment Comparison</a></li>
<li><a href="#support" id="toc-support" class="nav-link" data-scroll-target="#support">Support</a></li>
<li><a href="#license" id="toc-license" class="nav-link" data-scroll-target="#license">License</a></li>
<li><a href="#acknowledgements" id="toc-acknowledgements" class="nav-link" data-scroll-target="#acknowledgements">Acknowledgements</a></li>
</ul></li>
<li><a href="#adding-a-new-integration" id="toc-adding-a-new-integration" class="nav-link" data-scroll-target="#adding-a-new-integration">Adding a new integration</a></li>
</ul>
</nav>
</div>
<!-- main -->
<main class="content" id="quarto-document-content">
<header id="title-block-header" class="quarto-title-block default"><nav class="quarto-page-breadcrumbs quarto-title-breadcrumbs d-none d-lg-block" aria-label="breadcrumb"><ol class="breadcrumb"><li class="breadcrumb-item"><a href="../docs/fsdp_qlora.html">Advanced Features</a></li><li class="breadcrumb-item"><a href="../docs/custom_integrations.html">Custom Integrations</a></li></ol></nav>
<div class="quarto-title">
<h1 class="title">Custom Integrations</h1>
</div>
<div class="quarto-title-meta">
</div>
</header>
<p>Axolotl adds custom features through <code>integrations</code>. They are located within the <code>src/axolotl/integrations</code> directory.</p>
<p>To enable them, please check the respective documentations.</p>
<section id="cut-cross-entropy" class="level2">
<h2 class="anchored" data-anchor-id="cut-cross-entropy">Cut Cross Entropy</h2>
<p>Cut Cross Entropy (CCE) reduces VRAM usage through optimization on the cross-entropy operation during loss calculation.</p>
<p>See https://github.com/apple/ml-cross-entropy</p>
<section id="requirements" class="level3">
<h3 class="anchored" data-anchor-id="requirements">Requirements</h3>
<ul>
<li>PyTorch 2.4.0 or higher</li>
</ul>
</section>
<section id="installation" class="level3">
<h3 class="anchored" data-anchor-id="installation">Installation</h3>
<p>Run the following command to install <code>cut_cross_entropy[transformers]</code> if you don’t have it already.</p>
<ul>
<li>If you are in dev environment</li>
</ul>
<div class="code-copy-outer-scaffold"><div class="sourceCode" id="cb1"><pre class="sourceCode bash code-with-copy"><code class="sourceCode bash"><span id="cb1-1"><a href="#cb1-1" aria-hidden="true" tabindex="-1"></a><span class="ex">python</span> scripts/cutcrossentropy_install.py <span class="kw">|</span> <span class="fu">sh</span></span></code></pre></div><button title="Copy to Clipboard" class="code-copy-button"><i class="bi"></i></button></div>
<ul>
<li>If you are installing from pip</li>
</ul>
<div class="code-copy-outer-scaffold"><div class="sourceCode" id="cb2"><pre class="sourceCode bash code-with-copy"><code class="sourceCode bash"><span id="cb2-1"><a href="#cb2-1" aria-hidden="true" tabindex="-1"></a><span class="ex">pip3</span> uninstall <span class="at">-y</span> cut-cross-entropy <span class="kw">&amp;&amp;</span> <span class="ex">pip3</span> install <span class="st">"cut-cross-entropy[transformers] @ git+https://github.com/axolotl-ai-cloud/ml-cross-entropy.git@fec1a88"</span></span></code></pre></div><button title="Copy to Clipboard" class="code-copy-button"><i class="bi"></i></button></div>
</section>
<section id="usage" class="level3">
<h3 class="anchored" data-anchor-id="usage">Usage</h3>
<div class="code-copy-outer-scaffold"><div class="sourceCode" id="cb3"><pre class="sourceCode yaml code-with-copy"><code class="sourceCode yaml"><span id="cb3-1"><a href="#cb3-1" aria-hidden="true" tabindex="-1"></a><span class="fu">plugins</span><span class="kw">:</span></span>
<span id="cb3-2"><a href="#cb3-2" aria-hidden="true" tabindex="-1"></a><span class="at"> </span><span class="kw">-</span><span class="at"> axolotl.integrations.cut_cross_entropy.CutCrossEntropyPlugin</span></span></code></pre></div><button title="Copy to Clipboard" class="code-copy-button"><i class="bi"></i></button></div>
</section>
<section id="supported-models" class="level3">
<h3 class="anchored" data-anchor-id="supported-models">Supported Models</h3>
<ul>
<li>afmoe</li>
<li>apertus</li>
<li>arcee</li>
<li>cohere</li>
<li>cohere2</li>
<li>deepseek_v3</li>
<li>exaone4</li>
<li>gemma</li>
<li>gemma2</li>
<li>gemma3</li>
<li>gemma3_text</li>
<li>gemma3n</li>
<li>gemma3n_text</li>
<li>gemma4</li>
<li>glm</li>
<li>glm4</li>
<li>glm4_moe</li>
<li>glm4_moe_lite</li>
<li>glm46v</li>
<li>glm4v</li>
<li>glm4v_moe</li>
<li>glm_image</li>
<li>glm_moe_dsa</li>
<li>gpt_oss</li>
<li>granite</li>
<li>granitemoe</li>
<li>granitemoehybrid</li>
<li>granitemoeshared</li>
<li>hunyuan_v1_dense</li>
<li>hunyuan_v1_moe</li>
<li>internvl</li>
<li>kimi_linear</li>
<li>lfm2</li>
<li>lfm2_moe</li>
<li>lfm2_vl</li>
<li>llama</li>
<li>llama4</li>
<li>llama4_text</li>
<li>llava</li>
<li>ministral</li>
<li>ministral3</li>
<li>mistral</li>
<li>mistral3</li>
<li>mistral4</li>
<li>mixtral</li>
<li>mllama</li>
<li>nemotron_h</li>
<li>olmo</li>
<li>olmo2</li>
<li>olmo3</li>
<li>olmoe</li>
<li>phi</li>
<li>phi3</li>
<li>phi4_multimodal</li>
<li>qwen2</li>
<li>qwen2_5_vl</li>
<li>qwen2_moe</li>
<li>qwen2_vl</li>
<li>qwen3</li>
<li>qwen3_5</li>
<li>qwen3_5_text</li>
<li>qwen3_5_moe</li>
<li>qwen3_5_moe_text</li>
<li>qwen3_moe</li>
<li>qwen3_next</li>
<li>qwen3_vl</li>
<li>qwen3_vl_moe</li>
<li>seed_oss</li>
<li>smollm3</li>
<li>step3p5</li>
<li>voxtral</li>
</ul>
</section>
<section id="citation" class="level3">
<h3 class="anchored" data-anchor-id="citation">Citation</h3>
<div class="code-copy-outer-scaffold"><div class="sourceCode" id="cb4"><pre class="sourceCode bib code-with-copy"><code class="sourceCode bibtex"><span id="cb4-1"><a href="#cb4-1" aria-hidden="true" tabindex="-1"></a><span class="va">@article</span>{<span class="ot">wijmans2024cut</span>,</span>
<span id="cb4-2"><a href="#cb4-2" aria-hidden="true" tabindex="-1"></a> <span class="dt">author</span> = {Erik Wijmans and</span>
<span id="cb4-3"><a href="#cb4-3" aria-hidden="true" tabindex="-1"></a> Brody Huval and</span>
<span id="cb4-4"><a href="#cb4-4" aria-hidden="true" tabindex="-1"></a> Alexander Hertzberg and</span>
<span id="cb4-5"><a href="#cb4-5" aria-hidden="true" tabindex="-1"></a> Vladlen Koltun and</span>
<span id="cb4-6"><a href="#cb4-6" aria-hidden="true" tabindex="-1"></a> Philipp Kr<span class="ch">\"</span>ahenb<span class="ch">\"</span>uhl},</span>
<span id="cb4-7"><a href="#cb4-7" aria-hidden="true" tabindex="-1"></a> <span class="dt">title</span> = {Cut Your Losses in Large-Vocabulary Language Models},</span>
<span id="cb4-8"><a href="#cb4-8" aria-hidden="true" tabindex="-1"></a> <span class="dt">journal</span> = {arXiv},</span>
<span id="cb4-9"><a href="#cb4-9" aria-hidden="true" tabindex="-1"></a> <span class="dt">year</span> = {2024},</span>
<span id="cb4-10"><a href="#cb4-10" aria-hidden="true" tabindex="-1"></a> <span class="dt">url</span> = {https://arxiv.org/abs/2411.09009},</span>
<span id="cb4-11"><a href="#cb4-11" aria-hidden="true" tabindex="-1"></a>}</span></code></pre></div><button title="Copy to Clipboard" class="code-copy-button"><i class="bi"></i></button></div>
<p>Please see reference <a href="https://github.com/axolotl-ai-cloud/axolotl/tree/main/src/axolotl/integrations/cut_cross_entropy">here</a></p>
</section>
</section>
<section id="densemixer" class="level2">
<h2 class="anchored" data-anchor-id="densemixer">DenseMixer</h2>
<p>See <a href="https://github.com/yaof20/DenseMixer/">DenseMixer</a></p>
<p>Simply add the following to your axolotl YAML config:</p>
<div class="code-copy-outer-scaffold"><div class="sourceCode" id="cb5"><pre class="sourceCode yaml code-with-copy"><code class="sourceCode yaml"><span id="cb5-1"><a href="#cb5-1" aria-hidden="true" tabindex="-1"></a><span class="fu">plugins</span><span class="kw">:</span></span>
<span id="cb5-2"><a href="#cb5-2" aria-hidden="true" tabindex="-1"></a><span class="at"> </span><span class="kw">-</span><span class="at"> axolotl.integrations.densemixer.DenseMixerPlugin</span></span></code></pre></div><button title="Copy to Clipboard" class="code-copy-button"><i class="bi"></i></button></div>
<p>Please see reference <a href="https://github.com/axolotl-ai-cloud/axolotl/tree/main/src/axolotl/integrations/densemixer">here</a></p>
</section>
<section id="diffusion-lm-training-plugin-for-axolotl" class="level2">
<h2 class="anchored" data-anchor-id="diffusion-lm-training-plugin-for-axolotl">Diffusion LM Training Plugin for Axolotl</h2>
<p>This plugin enables diffusion language model training using an approach inspired by
LLaDA (Large Language Diffusion Models) within Axolotl.</p>
<section id="overview" class="level3">
<h3 class="anchored" data-anchor-id="overview">Overview</h3>
<p>LLaDA is a diffusion-based approach to language model training that uses:
- <strong>Random token masking</strong> during training instead of next-token prediction
- <strong>Bidirectional attention</strong> to allow the model to attend to the full context
- <strong>Importance weighting</strong> based on masking probabilities for stable training</p>
<p>This approach can lead to more robust language models with better understanding of
bidirectional context.</p>
</section>
<section id="installation-1" class="level3">
<h3 class="anchored" data-anchor-id="installation-1">Installation</h3>
<p>The plugin is included with Axolotl. See our
<a href="https://docs.axolotl.ai/docs/installation.html">installation docs</a>.</p>
</section>
<section id="quickstart" class="level3">
<h3 class="anchored" data-anchor-id="quickstart">Quickstart</h3>
<p>Train with an example config (Llama‑3.2 1B):
- Pretrain: <code>axolotl train examples/llama-3/diffusion-3.2-1b-pretrain.yaml</code>
- SFT: <code>axolotl train examples/llama-3/diffusion-3.2-1b-sft.yaml</code></p>
</section>
<section id="basic-configuration" class="level3">
<h3 class="anchored" data-anchor-id="basic-configuration">Basic Configuration</h3>
<p>You can also modify your existing configs to enable / customize diffusion training.</p>
<p>Add the following to your Axolotl config:</p>
<div class="code-copy-outer-scaffold"><div class="sourceCode" id="cb6"><pre class="sourceCode yaml code-with-copy"><code class="sourceCode yaml"><span id="cb6-1"><a href="#cb6-1" aria-hidden="true" tabindex="-1"></a><span class="fu">plugins</span><span class="kw">:</span></span>
<span id="cb6-2"><a href="#cb6-2" aria-hidden="true" tabindex="-1"></a><span class="at"> </span><span class="kw">-</span><span class="at"> axolotl.integrations.diffusion.DiffusionPlugin</span></span></code></pre></div><button title="Copy to Clipboard" class="code-copy-button"><i class="bi"></i></button></div>
<p>And, configure the nested <code>diffusion</code> block (defaults shown):</p>
<div class="code-copy-outer-scaffold"><div class="sourceCode" id="cb7"><pre class="sourceCode yaml code-with-copy"><code class="sourceCode yaml"><span id="cb7-1"><a href="#cb7-1" aria-hidden="true" tabindex="-1"></a><span class="fu">diffusion</span><span class="kw">:</span></span>
<span id="cb7-2"><a href="#cb7-2" aria-hidden="true" tabindex="-1"></a><span class="at"> </span><span class="fu">noise_schedule</span><span class="kw">:</span><span class="at"> linear</span><span class="co"> # or "cosine"</span></span>
<span id="cb7-3"><a href="#cb7-3" aria-hidden="true" tabindex="-1"></a><span class="at"> </span><span class="fu">min_mask_ratio</span><span class="kw">:</span><span class="at"> </span><span class="fl">0.1</span></span>
<span id="cb7-4"><a href="#cb7-4" aria-hidden="true" tabindex="-1"></a><span class="at"> </span><span class="fu">max_mask_ratio</span><span class="kw">:</span><span class="at"> </span><span class="fl">0.9</span></span>
<span id="cb7-5"><a href="#cb7-5" aria-hidden="true" tabindex="-1"></a><span class="at"> </span><span class="fu">num_diffusion_steps</span><span class="kw">:</span><span class="at"> </span><span class="dv">128</span></span>
<span id="cb7-6"><a href="#cb7-6" aria-hidden="true" tabindex="-1"></a><span class="at"> </span><span class="fu">eps</span><span class="kw">:</span><span class="at"> </span><span class="fl">1e-3</span></span>
<span id="cb7-7"><a href="#cb7-7" aria-hidden="true" tabindex="-1"></a><span class="at"> </span><span class="fu">importance_weighting</span><span class="kw">:</span><span class="at"> </span><span class="ch">true</span></span>
<span id="cb7-8"><a href="#cb7-8" aria-hidden="true" tabindex="-1"></a></span>
<span id="cb7-9"><a href="#cb7-9" aria-hidden="true" tabindex="-1"></a><span class="co"> # Mask token (training auto-adds if missing, avoid pad/eos)</span></span>
<span id="cb7-10"><a href="#cb7-10" aria-hidden="true" tabindex="-1"></a><span class="at"> </span><span class="fu">mask_token_str</span><span class="kw">:</span><span class="at"> </span><span class="st">"&lt;|diffusion_mask|&gt;"</span></span>
<span id="cb7-11"><a href="#cb7-11" aria-hidden="true" tabindex="-1"></a><span class="co"> # Or use an existing special token id (e.g., 128002 for Llama-3.x)</span></span>
<span id="cb7-12"><a href="#cb7-12" aria-hidden="true" tabindex="-1"></a><span class="co"> # mask_token_id: 128002</span></span>
<span id="cb7-13"><a href="#cb7-13" aria-hidden="true" tabindex="-1"></a></span>
<span id="cb7-14"><a href="#cb7-14" aria-hidden="true" tabindex="-1"></a><span class="co"> # Sample generation during training (optional)</span></span>
<span id="cb7-15"><a href="#cb7-15" aria-hidden="true" tabindex="-1"></a><span class="at"> </span><span class="fu">generate_samples</span><span class="kw">:</span><span class="at"> </span><span class="ch">true</span></span>
<span id="cb7-16"><a href="#cb7-16" aria-hidden="true" tabindex="-1"></a><span class="at"> </span><span class="fu">generation_interval</span><span class="kw">:</span><span class="at"> </span><span class="dv">100</span></span>
<span id="cb7-17"><a href="#cb7-17" aria-hidden="true" tabindex="-1"></a><span class="at"> </span><span class="fu">num_generation_samples</span><span class="kw">:</span><span class="at"> </span><span class="dv">3</span></span>
<span id="cb7-18"><a href="#cb7-18" aria-hidden="true" tabindex="-1"></a><span class="at"> </span><span class="fu">generation_steps</span><span class="kw">:</span><span class="at"> </span><span class="dv">128</span></span>
<span id="cb7-19"><a href="#cb7-19" aria-hidden="true" tabindex="-1"></a><span class="at"> </span><span class="fu">generation_temperature</span><span class="kw">:</span><span class="at"> </span><span class="fl">0.0</span></span>
<span id="cb7-20"><a href="#cb7-20" aria-hidden="true" tabindex="-1"></a><span class="at"> </span><span class="fu">generation_max_length</span><span class="kw">:</span><span class="at"> </span><span class="dv">100</span></span></code></pre></div><button title="Copy to Clipboard" class="code-copy-button"><i class="bi"></i></button></div>
</section>
<section id="supported-models-1" class="level3">
<h3 class="anchored" data-anchor-id="supported-models-1">Supported Models</h3>
<p>Any models that support 4D attention masks should work out of the box. If not, please
create an <a href="https://github.com/axolotl-ai-cloud/axolotl/issues">issue</a> or open a
<a href="https://github.com/axolotl-ai-cloud/axolotl/compare">PR</a>!</p>
</section>
<section id="how-it-works" class="level3">
<h3 class="anchored" data-anchor-id="how-it-works">How It Works</h3>
</section>
<section id="random-masking" class="level3">
<h3 class="anchored" data-anchor-id="random-masking">Random Masking</h3>
<p>During training, tokens are randomly masked:
- Sample timestep <code>t</code> uniformly from [0, 1]
- Calculate masking probability: <code>p = (1 - eps) * t + eps</code>
- Randomly mask tokens with probability <code>p</code></p>
</section>
<section id="diffusion-loss" class="level3">
<h3 class="anchored" data-anchor-id="diffusion-loss">Diffusion Loss</h3>
<p>Loss is computed only on masked tokens with (optional) importance weighting:</p>
<div class="code-copy-outer-scaffold"><div class="sourceCode" id="cb8"><pre class="sourceCode python code-with-copy"><code class="sourceCode python"><span id="cb8-1"><a href="#cb8-1" aria-hidden="true" tabindex="-1"></a>loss <span class="op">=</span> <span class="bu">sum</span>(cross_entropy(pred, target) <span class="op">/</span> p_mask) <span class="op">/</span> total_tokens</span></code></pre></div><button title="Copy to Clipboard" class="code-copy-button"><i class="bi"></i></button></div>
</section>
<section id="sample-generation" class="level3">
<h3 class="anchored" data-anchor-id="sample-generation">Sample Generation</h3>
<p>When <code>diffusion.generate_samples: true</code>, the plugin generates samples during training:</p>
<pre><code>Sample 1:
Original (45 tokens): The quick brown fox jumps over the lazy dog...
Masked (18/45 tokens, 40.0%): The [MASK] [MASK] fox [MASK] over [MASK] lazy [MASK]...
Generated: The quick brown fox jumps over the lazy dog...</code></pre>
<p>Samples are logged to console and wandb (if enabled).</p>
</section>
<section id="inference" class="level3">
<h3 class="anchored" data-anchor-id="inference">Inference</h3>
<p>Diffusion inference is integrated into the standard Axolotl CLI. Use the same config
you trained with and run:</p>
<pre><code>axolotl inference path/to/your-config.yaml</code></pre>
<p>Optionally, pass <code>--gradio</code> to use a simple web interface.</p>
<p>Interactive controls (prefix the prompt with commands):
- <code>:complete N</code> → completion mode with N new masked tokens appended (default 64)
- <code>:mask R</code> → random masking mode with target mask ratio R in [0.0, 1.0]</p>
<p>Example session:</p>
<pre><code>================================================================================
Commands:
:complete N -&gt; completion mode with N tokens (default 64)
:mask R -&gt; random masking with ratio R (0.0–1.0)
================================================================================
Give me an instruction (Ctrl + D to submit):
:mask 0.4 The quick brown fox jumps over the lazy dog
Masked (40.0%):
The [MASK] brown [MASK] jumps over the [MASK] dog
Generated:
The quick brown fox jumps over the loud dog</code></pre>
</section>
<section id="metrics-and-monitoring" class="level3">
<h3 class="anchored" data-anchor-id="metrics-and-monitoring">Metrics and Monitoring</h3>
<p>The plugin adds (or modifies) several metrics to track diffusion training:</p>
<ul>
<li><code>train/loss</code>: Weighted diffusion loss</li>
<li><code>train/accuracy</code>: Accuracy on masked tokens</li>
<li><code>train/mask_ratio</code>: Average fraction of tokens masked</li>
<li><code>train/num_masked_tokens</code>: Number of tokens masked</li>
<li><code>train/avg_p_mask</code>: Average masking probability</li>
<li><code>train/ce_loss</code>: Unweighted cross-entropy loss</li>
<li><code>train/importance_weight_avg</code>: Average importance weight</li>
</ul>
</section>
<section id="limitations" class="level3">
<h3 class="anchored" data-anchor-id="limitations">Limitations</h3>
<ul>
<li>No flash attention support</li>
<li>No RL training support</li>
</ul>
</section>
<section id="references" class="level3">
<h3 class="anchored" data-anchor-id="references">References</h3>
<ul>
<li><a href="https://arxiv.org/abs/2404.10406">LLaDA Paper</a></li>
<li><a href="https://docs.axolotl.ai/">Axolotl Documentation</a></li>
<li><a href="https://docs.axolotl.ai/docs/api/integrations.diffusion.args.html#axolotl.integrations.diffusion.args">API reference for plugin</a></li>
</ul>
<p>Please see reference <a href="https://github.com/axolotl-ai-cloud/axolotl/tree/main/src/axolotl/integrations/diffusion">here</a></p>
</section>
</section>
<section id="grokfast" class="level2">
<h2 class="anchored" data-anchor-id="grokfast">Grokfast</h2>
<p>See https://github.com/ironjr/grokfast</p>
<section id="usage-1" class="level3">
<h3 class="anchored" data-anchor-id="usage-1">Usage</h3>
<div class="code-copy-outer-scaffold"><div class="sourceCode" id="cb12"><pre class="sourceCode yaml code-with-copy"><code class="sourceCode yaml"><span id="cb12-1"><a href="#cb12-1" aria-hidden="true" tabindex="-1"></a><span class="fu">plugins</span><span class="kw">:</span></span>
<span id="cb12-2"><a href="#cb12-2" aria-hidden="true" tabindex="-1"></a><span class="at"> </span><span class="kw">-</span><span class="at"> axolotl.integrations.grokfast.GrokfastPlugin</span></span>
<span id="cb12-3"><a href="#cb12-3" aria-hidden="true" tabindex="-1"></a></span>
<span id="cb12-4"><a href="#cb12-4" aria-hidden="true" tabindex="-1"></a><span class="fu">grokfast_alpha</span><span class="kw">:</span><span class="at"> </span><span class="fl">2.0</span></span>
<span id="cb12-5"><a href="#cb12-5" aria-hidden="true" tabindex="-1"></a><span class="fu">grokfast_lamb</span><span class="kw">:</span><span class="at"> </span><span class="fl">0.98</span></span></code></pre></div><button title="Copy to Clipboard" class="code-copy-button"><i class="bi"></i></button></div>
</section>
<section id="citation-1" class="level3">
<h3 class="anchored" data-anchor-id="citation-1">Citation</h3>
<div class="code-copy-outer-scaffold"><div class="sourceCode" id="cb13"><pre class="sourceCode bib code-with-copy"><code class="sourceCode bibtex"><span id="cb13-1"><a href="#cb13-1" aria-hidden="true" tabindex="-1"></a><span class="va">@article</span>{<span class="ot">lee2024grokfast</span>,</span>
<span id="cb13-2"><a href="#cb13-2" aria-hidden="true" tabindex="-1"></a> <span class="dt">title</span>={{Grokfast}: Accelerated Grokking by Amplifying Slow Gradients},</span>
<span id="cb13-3"><a href="#cb13-3" aria-hidden="true" tabindex="-1"></a> <span class="dt">author</span>={Lee, Jaerin and Kang, Bong Gyun and Kim, Kihoon and Lee, Kyoung Mu},</span>
<span id="cb13-4"><a href="#cb13-4" aria-hidden="true" tabindex="-1"></a> <span class="dt">journal</span>={arXiv preprint arXiv:2405.20233},</span>
<span id="cb13-5"><a href="#cb13-5" aria-hidden="true" tabindex="-1"></a> <span class="dt">year</span>={2024}</span>
<span id="cb13-6"><a href="#cb13-6" aria-hidden="true" tabindex="-1"></a>}</span></code></pre></div><button title="Copy to Clipboard" class="code-copy-button"><i class="bi"></i></button></div>
<p>Please see reference <a href="https://github.com/axolotl-ai-cloud/axolotl/tree/main/src/axolotl/integrations/grokfast">here</a></p>
</section>
</section>
<section id="kernels-integration" class="level2">
<h2 class="anchored" data-anchor-id="kernels-integration">Kernels Integration</h2>
<p>MoE (Mixture of Experts) kernels speed up training for MoE layers and reduce VRAM costs. In transformers v5, <code>batched_mm</code> and <code>grouped_mm</code> were integrated as built-in options via the <code>experts_implementation</code> config kwarg:</p>
<div class="code-copy-outer-scaffold"><div class="sourceCode" id="cb14"><pre class="sourceCode python code-with-copy"><code class="sourceCode python"><span id="cb14-1"><a href="#cb14-1" aria-hidden="true" tabindex="-1"></a><span class="kw">class</span> ExpertsInterface(GeneralInterface):</span>
<span id="cb14-2"><a href="#cb14-2" aria-hidden="true" tabindex="-1"></a> _global_mapping <span class="op">=</span> {</span>
<span id="cb14-3"><a href="#cb14-3" aria-hidden="true" tabindex="-1"></a> <span class="st">"batched_mm"</span>: batched_mm_experts_forward,</span>
<span id="cb14-4"><a href="#cb14-4" aria-hidden="true" tabindex="-1"></a> <span class="st">"grouped_mm"</span>: grouped_mm_experts_forward,</span>
<span id="cb14-5"><a href="#cb14-5" aria-hidden="true" tabindex="-1"></a> }</span></code></pre></div><button title="Copy to Clipboard" class="code-copy-button"><i class="bi"></i></button></div>
<p>In our custom integration, we add support for <strong>ScatterMoE</strong> and <strong>SonicMoE</strong>, which are more efficient and faster than <code>grouped_mm</code>.</p>
<section id="usage-2" class="level3">
<h3 class="anchored" data-anchor-id="usage-2">Usage</h3>
<p>Add the following to your axolotl YAML config:</p>
<div class="code-copy-outer-scaffold"><div class="sourceCode" id="cb15"><pre class="sourceCode yaml code-with-copy"><code class="sourceCode yaml"><span id="cb15-1"><a href="#cb15-1" aria-hidden="true" tabindex="-1"></a><span class="fu">plugins</span><span class="kw">:</span></span>
<span id="cb15-2"><a href="#cb15-2" aria-hidden="true" tabindex="-1"></a><span class="at"> </span><span class="kw">-</span><span class="at"> axolotl.integrations.kernels.KernelsPlugin</span></span>
<span id="cb15-3"><a href="#cb15-3" aria-hidden="true" tabindex="-1"></a></span>
<span id="cb15-4"><a href="#cb15-4" aria-hidden="true" tabindex="-1"></a><span class="fu">use_kernels</span><span class="kw">:</span><span class="at"> </span><span class="ch">true</span></span>
<span id="cb15-5"><a href="#cb15-5" aria-hidden="true" tabindex="-1"></a></span>
<span id="cb15-6"><a href="#cb15-6" aria-hidden="true" tabindex="-1"></a><span class="fu">use_scattermoe</span><span class="kw">:</span><span class="at"> </span><span class="ch">true</span></span>
<span id="cb15-7"><a href="#cb15-7" aria-hidden="true" tabindex="-1"></a><span class="fu">use_sonicmoe</span><span class="kw">:</span><span class="at"> </span><span class="ch">true</span></span></code></pre></div><button title="Copy to Clipboard" class="code-copy-button"><i class="bi"></i></button></div>
<p><strong>Important:</strong> Setting <code>experts_implementation</code> to <code>batched_mm</code> or <code>grouped_mm</code> is incompatible with custom kernel options. The exception is <code>experts_implementation: scattermoe</code>, which is used for models like Gemma 4 that embed MoE directly in the decoder layer (no SparseMoeBlock) and dispatch through the transformers <code>ExpertsInterface</code>.</p>
</section>
<section id="sonicmoe-installation" class="level3">
<h3 class="anchored" data-anchor-id="sonicmoe-installation">SonicMoE installation</h3>
<p><strong>Prerequisites:</strong>
- NVIDIA Hopper (H100, H200) or Blackwell (B200, GB200) GPU
- CUDA 12.9+ (13.0+ for B300)
- PyTorch 2.7+ (2.9.1 recommended)
- For B300: Triton 3.6.0</p>
<div class="code-copy-outer-scaffold"><div class="sourceCode" id="cb16"><pre class="sourceCode bash code-with-copy"><code class="sourceCode bash"><span id="cb16-1"><a href="#cb16-1" aria-hidden="true" tabindex="-1"></a><span class="ex">pip</span> install <span class="at">--ignore-requires-python</span> <span class="at">--no-deps</span> <span class="st">"sonic-moe @ git+https://github.com/Dao-AILab/sonic-moe.git@116e2df0a41874f77fa0ad269ce7df3f0cfcb956"</span> <span class="kw">&amp;&amp;</span> <span class="ex">pip</span> install nvidia-cutlass-dsl==4.4.0 quack-kernels==0.2.5</span></code></pre></div><button title="Copy to Clipboard" class="code-copy-button"><i class="bi"></i></button></div>
<p>See the <a href="https://github.com/Dao-AILab/sonic-moe?tab=readme-ov-file#-installation">SonicMoE installation guide</a> for the latest prerequisite details.</p>
<p><strong>Note:</strong> Blackwell support is in upstream beta. On Blackwell GPUs, Axolotl automatically sets <code>USE_QUACK_GEMM=1</code> to enable the Blackwell kernels.</p>
</section>
<section id="how-it-works-1" class="level3">
<h3 class="anchored" data-anchor-id="how-it-works-1">How It Works</h3>
<p>The <code>KernelsPlugin</code> runs before model loading and:</p>
</section>
<section id="scattermoe" class="level3">
<h3 class="anchored" data-anchor-id="scattermoe">ScatterMoE</h3>
<ol type="1">
<li>Registers the ScatterMoE kernel from the local <code>libs/scattermoe_lora</code> package (includes fused LoRA support via Triton kernels).</li>
<li>Patches the model’s <code>SparseMoeBlock</code> forward method with the optimized ScatterMoE implementation via the HF <code>kernels</code> library.</li>
</ol>
</section>
<section id="sonicmoe" class="level3">
<h3 class="anchored" data-anchor-id="sonicmoe">SonicMoE</h3>
<ol type="1">
<li>Resolves the model’s MoE block class(es) from <code>constants.py</code>.</li>
<li>Patches the forward method with SonicMoE’s optimized CUTLASS kernels and registers a weight converter for the interleaved gate/up projection format.</li>
<li>Supports pluggable routing strategies (see routing table below).</li>
</ol>
<p>Both paths use the shared <code>resolve_moe_block_classes</code> utility in <code>constants.py</code> for model-type-to-class resolution.</p>
</section>
<section id="model-support-matrix" class="level3">
<h3 class="anchored" data-anchor-id="model-support-matrix">Model Support Matrix</h3>
<p>Most models use the <strong>SwiGLU</strong> activation (<code>silu(gate) * up</code>). Gemma 4 uses <strong>GEGLU</strong> (<code>gelu(gate) * up</code>). ScatterMoE supports any gated activation (activation is applied in Python between kernel calls). SonicMoE supports SwiGLU, GEGLU, and REGLU via its <code>ActivationType</code> enum.</p>
</section>
<section id="routing-strategies" class="level3">
<h3 class="anchored" data-anchor-id="routing-strategies">Routing strategies</h3>
<table class="caption-top table">
<colgroup>
<col style="width: 18%">
<col style="width: 18%">
<col style="width: 31%">
<col style="width: 31%">
</colgroup>
<thead>
<tr class="header">
<th>Routing Strategy</th>
<th>Description</th>
<th style="text-align: center;">ScatterMoE</th>
<th style="text-align: center;">SonicMoE</th>
</tr>
</thead>
<tbody>
<tr class="odd">
<td>softmax → topk</td>
<td>Softmax over experts, select top-K, optional renormalization</td>
<td style="text-align: center;">Yes</td>
<td style="text-align: center;">Yes</td>
</tr>
<tr class="even">
<td>softmax → group selection → topk</td>
<td>Softmax, select top groups (sum of top-2 per group), topk from selected groups, renorm + scaling</td>
<td style="text-align: center;">No</td>
<td style="text-align: center;">Yes</td>
</tr>
<tr class="odd">
<td>sigmoid → topk (with groups)</td>
<td>Sigmoid + bias correction, group-based masking, topk from masked scores, weights from original sigmoid</td>
<td style="text-align: center;">Yes</td>
<td style="text-align: center;">Yes</td>
</tr>
<tr class="even">
<td>sigmoid → topk (no groups)</td>
<td>Sigmoid + bias correction, straight topk (n_group=1)</td>
<td style="text-align: center;">Yes</td>
<td style="text-align: center;">Yes</td>
</tr>
<tr class="odd">
<td>softmax → bias correction → topk</td>
<td>Softmax, bias via <code>gate.moe_statics</code>, topk, gather from original probs, clamp-based renorm</td>
<td style="text-align: center;">No</td>
<td style="text-align: center;">Yes</td>
</tr>
<tr class="even">
<td>softmax → group_limited_greedy</td>
<td>Softmax, group selection (max per group), topk, scale only (no renorm)</td>
<td style="text-align: center;">No</td>
<td style="text-align: center;">Yes</td>
</tr>
<tr class="odd">
<td>softmax → topk via gate.wg</td>
<td>Softmax, gate weight at <code>gate.wg.weight</code> (not <code>gate.weight</code>), always renormalize</td>
<td style="text-align: center;">No</td>
<td style="text-align: center;">Yes</td>
</tr>
<tr class="even">
<td>softmax → topk + per_expert_scale</td>
<td>RMSNorm → scale → proj → softmax → topk → renorm → per-expert learned scales</td>
<td style="text-align: center;">Yes</td>
<td style="text-align: center;">Yes</td>
</tr>
<tr class="odd">
<td>fused topk → softmax</td>
<td>Routing + expert computation fused in a single kernel</td>
<td style="text-align: center;">No</td>
<td style="text-align: center;">Planned</td>
</tr>
</tbody>
</table>
</section>
<section id="per-model-support" class="level3">
<h3 class="anchored" data-anchor-id="per-model-support">Per-model support</h3>
<table class="caption-top table">
<colgroup>
<col style="width: 15%">
<col style="width: 15%">
<col style="width: 15%">
<col style="width: 26%">
<col style="width: 26%">
</colgroup>
<thead>
<tr class="header">
<th>Model Type</th>
<th>Architecture</th>
<th>Routing</th>
<th style="text-align: center;">ScatterMoE</th>
<th style="text-align: center;">SonicMoE</th>
</tr>
</thead>
<tbody>
<tr class="odd">
<td><code>qwen2_moe</code></td>
<td>Qwen2-MoE</td>
<td>softmax → topk</td>
<td style="text-align: center;"><strong>Yes</strong></td>
<td style="text-align: center;"><strong>Yes</strong></td>
</tr>
<tr class="even">
<td><code>qwen3_moe</code></td>
<td>Qwen3-MoE</td>
<td>softmax → topk</td>
<td style="text-align: center;"><strong>Yes</strong></td>
<td style="text-align: center;"><strong>Yes</strong></td>
</tr>
<tr class="odd">
<td><code>qwen3_5_moe</code></td>
<td>Qwen3.5-MoE</td>
<td>softmax → topk</td>
<td style="text-align: center;"><strong>Yes</strong></td>
<td style="text-align: center;"><strong>Yes</strong></td>
</tr>
<tr class="even">
<td><code>qwen3_5_moe_text</code></td>
<td>Qwen3.5-MoE (VLM text)</td>
<td>softmax → topk</td>
<td style="text-align: center;"><strong>Yes</strong></td>
<td style="text-align: center;"><strong>Yes</strong></td>
</tr>
<tr class="odd">
<td><code>qwen3_next</code></td>
<td>Qwen3-Next</td>
<td>softmax → topk</td>
<td style="text-align: center;"><strong>Yes</strong></td>
<td style="text-align: center;"><strong>Yes</strong></td>
</tr>
<tr class="even">
<td><code>qwen3_vl_moe</code></td>
<td>Qwen3-VL-MoE</td>
<td>softmax → topk</td>
<td style="text-align: center;"><strong>Yes</strong></td>
<td style="text-align: center;"><strong>Yes</strong></td>
</tr>
<tr class="odd">
<td><code>qwen3_omni_moe</code></td>
<td>Qwen3-Omni (Thinker + Talker)</td>
<td>softmax → topk</td>
<td style="text-align: center;"><strong>Yes</strong></td>
<td style="text-align: center;"><strong>Yes</strong></td>
</tr>
<tr class="even">
<td><code>olmoe</code></td>
<td>OLMoE</td>
<td>softmax → topk</td>
<td style="text-align: center;"><strong>Yes</strong></td>
<td style="text-align: center;"><strong>Yes</strong></td>
</tr>
<tr class="odd">
<td><code>mixtral</code></td>
<td>Mixtral</td>
<td>softmax → topk</td>
<td style="text-align: center;"><strong>Yes</strong></td>
<td style="text-align: center;"><strong>Yes</strong></td>
</tr>
<tr class="even">
<td><code>minimax</code></td>
<td>MiniMax</td>
<td>softmax → topk</td>
<td style="text-align: center;"><strong>Yes</strong></td>
<td style="text-align: center;"><strong>Yes</strong></td>
</tr>
<tr class="odd">
<td><code>mistral4</code></td>
<td>Mistral 4</td>
<td>softmax → group → topk</td>
<td style="text-align: center;">No</td>
<td style="text-align: center;"><strong>Yes</strong></td>
</tr>
<tr class="even">
<td><code>glm_moe_dsa</code></td>
<td>GLM-MoE DSA (GLM 5)</td>
<td>sigmoid → topk (groups)</td>
<td style="text-align: center;"><strong>Yes</strong></td>
<td style="text-align: center;"><strong>Yes</strong></td>
</tr>
<tr class="odd">
<td><code>deepseek_v3</code></td>
<td>DeepSeek-V3</td>
<td>sigmoid → topk (groups)</td>
<td style="text-align: center;"><strong>Yes</strong></td>
<td style="text-align: center;"><strong>Yes</strong></td>
</tr>
<tr class="even">
<td><code>glm4_moe</code></td>
<td>GLM4-MoE</td>
<td>sigmoid → topk (groups)</td>
<td style="text-align: center;"><strong>Yes</strong></td>
<td style="text-align: center;"><strong>Yes</strong></td>
</tr>
<tr class="odd">
<td><code>glm4_moe_lite</code></td>
<td>GLM4-MoE Lite (GLM 4.7 Flash)</td>
<td>sigmoid → topk (groups)</td>
<td style="text-align: center;"><strong>Yes</strong>*</td>
<td style="text-align: center;"><strong>Yes</strong></td>
</tr>
<tr class="even">
<td><code>glm4v_moe</code></td>
<td>GLM4v-MoE</td>
<td>sigmoid → topk (groups)</td>
<td style="text-align: center;"><strong>Yes</strong></td>
<td style="text-align: center;"><strong>Yes</strong></td>
</tr>
<tr class="odd">
<td><code>minimax_m2</code></td>
<td>MiniMax M2</td>
<td>sigmoid → topk (no groups)</td>
<td style="text-align: center;"><strong>Yes</strong></td>
<td style="text-align: center;"><strong>Yes</strong></td>
</tr>
<tr class="even">
<td><code>ernie4_5_moe</code></td>
<td>ERNIE 4.5 MoE</td>
<td>softmax → bias → topk</td>
<td style="text-align: center;">No</td>
<td style="text-align: center;"><strong>Yes</strong></td>
</tr>
<tr class="odd">
<td><code>deepseek_v2</code></td>
<td>DeepSeek-V2</td>
<td>softmax → group_limited_greedy</td>
<td style="text-align: center;">No</td>
<td style="text-align: center;"><strong>Yes</strong></td>
</tr>
<tr class="even">
<td><code>hunyuan_v1_moe</code></td>
<td>HunYuan V1 MoE</td>
<td>softmax → topk (gate.wg)</td>
<td style="text-align: center;">No</td>
<td style="text-align: center;"><strong>Yes</strong></td>
</tr>
<tr class="odd">
<td><code>gemma4_text</code></td>
<td>Gemma 4 (26B-A4B)</td>
<td>softmax → topk + per_expert_scale</td>
<td style="text-align: center;"><strong>Yes</strong>**</td>
<td style="text-align: center;"><strong>Yes</strong>**</td>
</tr>
<tr class="even">
<td><code>gpt_oss</code></td>
<td>GPT-OSS</td>
<td>fused topk → softmax</td>
<td style="text-align: center;">No</td>
<td style="text-align: center;">Planned</td>
</tr>
</tbody>
</table>
<p>* <code>glm4_moe_lite</code> with ScatterMoE may have issues — see Limitations.</p>
<p>** Gemma 4 uses <code>experts_implementation: scattermoe</code> path (registered via <code>ExpertsInterface</code>) instead of SparseMoeBlock patching, since Gemma 4 embeds MoE directly in its decoder layer (no separate SparseMoeBlock). See the <a href="#gemma-4">Gemma 4 section</a> below.</p>
</section>
<section id="feature-comparison" class="level3">
<h3 class="anchored" data-anchor-id="feature-comparison">Feature comparison</h3>
<table class="caption-top table">
<colgroup>
<col style="width: 23%">
<col style="width: 38%">
<col style="width: 38%">
</colgroup>
<thead>
<tr class="header">
<th>Feature</th>
<th style="text-align: center;">ScatterMoE</th>
<th style="text-align: center;">SonicMoE</th>
</tr>
</thead>
<tbody>
<tr class="odd">
<td>Kernel backend</td>
<td style="text-align: center;">Triton</td>
<td style="text-align: center;">CUTLASS</td>
</tr>
<tr class="even">
<td>GPU requirement</td>
<td style="text-align: center;">Any CUDA</td>
<td style="text-align: center;">Hopper (H100/H200) or Blackwell (B200+)</td>
</tr>
<tr class="odd">
<td>LoRA approach</td>
<td style="text-align: center;">Fused in Triton kernel</td>
<td style="text-align: center;">Runtime materialization + custom autograd</td>
</tr>
<tr class="even">
<td>LoRA overhead</td>
<td style="text-align: center;">Lower (fused computation)</td>
<td style="text-align: center;">Higher (per-forward materialization)</td>
</tr>
<tr class="odd">
<td>Gate/router LoRA</td>
<td style="text-align: center;">Yes</td>
<td style="text-align: center;">Yes</td>
</tr>
<tr class="even">
<td>Expert LoRA</td>
<td style="text-align: center;">Yes (fused)</td>
<td style="text-align: center;">Yes (materialized)</td>
</tr>
<tr class="odd">
<td>Shared expert LoRA</td>
<td style="text-align: center;">Yes (standard PEFT)</td>
<td style="text-align: center;">Yes (standard PEFT)</td>
</tr>
<tr class="even">
<td>Selective expert dequantization</td>
<td style="text-align: center;">Yes (~97% memory savings)</td>
<td style="text-align: center;">No</td>
</tr>
<tr class="odd">
<td>Weight format</td>
<td style="text-align: center;">Transposed <code>[E, hidden, 2*inter]</code></td>
<td style="text-align: center;">Interleaved gate/up <code>[2*I, H, E]</code></td>
</tr>
<tr class="even">
<td>torch.compile routing</td>
<td style="text-align: center;">No</td>
<td style="text-align: center;">Yes (optional)</td>
</tr>
</tbody>
</table>
</section>
<section id="shared-expert-handling" class="level3">
<h3 class="anchored" data-anchor-id="shared-expert-handling">Shared Expert Handling</h3>
<p>Both kernels handle shared experts identically. Shared expert attribute names are detected in order of priority:</p>
<ol type="1">
<li><code>shared_expert</code> (Qwen2-MoE)</li>
<li><code>shared_experts</code> (GLM-MoE, DeepSeek-V3)</li>
<li><code>shared_mlp</code> (HunYuan V1 MoE)</li>
</ol>
<p>If <code>shared_expert_gate</code> exists, sigmoid gating is applied to the shared expert contribution before adding it to the routed output. PEFT wraps shared expert linear layers with standard LoRA — no special handling is needed.</p>
</section>
<section id="gemma-4" class="level3">
<h3 class="anchored" data-anchor-id="gemma-4">Gemma 4</h3>
<p>Gemma 4 (e.g.&nbsp;<code>google/gemma-4-26B-A4B</code>) has a unique hybrid MoE architecture:</p>
<ul>
<li><strong>No SparseMoeBlock</strong>: MoE is embedded directly in the decoder layer alongside a dense MLP. Both run in parallel and their outputs are summed.</li>
<li><strong>Custom router</strong> (<code>Gemma4TextRouter</code>): RMSNorm → learned scale → linear projection → softmax → top-k → renormalization → per-expert learned scales.</li>
<li><strong>GEGLU activation</strong>: Uses <code>gelu_pytorch_tanh</code> (not SiLU/SwiGLU like most other MoE models).</li>
<li><strong>128 experts, top-k=8</strong> for the 26B-A4B variant.</li>
</ul>
<p>Because there is no SparseMoeBlock class to patch, Gemma 4 uses a different integration path: we register <code>"scattermoe"</code> as a custom implementation in the transformers <code>ExpertsInterface</code>, and set <code>experts_implementation: scattermoe</code> in the config. The <code>@use_experts_implementation</code> decorator on <code>Gemma4TextExperts</code> then dispatches to our ScatterMoE kernel automatically. The router is untouched — it runs as-is.</p>
</section>
<section id="limitations-1" class="level3">
<h3 class="anchored" data-anchor-id="limitations-1">Limitations</h3>
<ul>
<li><strong>ScatterMoE + GLM4-MoE Lite</strong>: ScatterMoE does not work reliably for GLM 4.7 Flash (<code>glm4_moe_lite</code>).</li>
<li><strong>Non-SwiGLU activations</strong>: Neither kernel supports MoE architectures with non-SwiGLU expert activations (e.g., GPT-OSS uses a custom GLU variant).</li>
<li><strong>GPT-OSS</strong>: Deferred — requires transposed weight layout <code>[E, H, 2*I]</code>, expert biases, and custom GLU activation. A dedicated forward path is needed.</li>
<li><strong>FSDP + fused gate LoRA (SonicMoE)</strong>: The fused topk→softmax path materializes a local tensor when LoRA delta is present to avoid DTensor + Tensor mixing under FSDP.</li>
</ul>
</section>
<section id="note-on-megablocks" class="level3">
<h3 class="anchored" data-anchor-id="note-on-megablocks">Note on MegaBlocks</h3>
<p>We tested <a href="https://huggingface.co/kernels-community/megablocks">MegaBlocks</a> but were unable to ensure numerical accuracy, so we did not integrate it. It was also incompatible with many newer model architectures in transformers.</p>
<p>Please see reference <a href="https://github.com/axolotl-ai-cloud/axolotl/tree/main/src/axolotl/integrations/kernels">here</a></p>
</section>
</section>
<section id="knowledge-distillation-kd" class="level2">
<h2 class="anchored" data-anchor-id="knowledge-distillation-kd">Knowledge Distillation (KD)</h2>
<section id="usage-3" class="level3">
<h3 class="anchored" data-anchor-id="usage-3">Usage</h3>
<div class="code-copy-outer-scaffold"><div class="sourceCode" id="cb17"><pre class="sourceCode yaml code-with-copy"><code class="sourceCode yaml"><span id="cb17-1"><a href="#cb17-1" aria-hidden="true" tabindex="-1"></a><span class="fu">plugins</span><span class="kw">:</span></span>
<span id="cb17-2"><a href="#cb17-2" aria-hidden="true" tabindex="-1"></a><span class="at"> </span><span class="kw">-</span><span class="at"> </span><span class="st">"axolotl.integrations.kd.KDPlugin"</span></span>
<span id="cb17-3"><a href="#cb17-3" aria-hidden="true" tabindex="-1"></a></span>
<span id="cb17-4"><a href="#cb17-4" aria-hidden="true" tabindex="-1"></a><span class="fu">kd_trainer</span><span class="kw">:</span><span class="at"> </span><span class="ch">True</span></span>
<span id="cb17-5"><a href="#cb17-5" aria-hidden="true" tabindex="-1"></a><span class="fu">kd_ce_alpha</span><span class="kw">:</span><span class="at"> </span><span class="fl">0.1</span></span>
<span id="cb17-6"><a href="#cb17-6" aria-hidden="true" tabindex="-1"></a><span class="fu">kd_alpha</span><span class="kw">:</span><span class="at"> </span><span class="fl">0.9</span></span>
<span id="cb17-7"><a href="#cb17-7" aria-hidden="true" tabindex="-1"></a><span class="fu">kd_temperature</span><span class="kw">:</span><span class="at"> </span><span class="fl">1.0</span></span>
<span id="cb17-8"><a href="#cb17-8" aria-hidden="true" tabindex="-1"></a></span>
<span id="cb17-9"><a href="#cb17-9" aria-hidden="true" tabindex="-1"></a><span class="fu">torch_compile</span><span class="kw">:</span><span class="at"> </span><span class="ch">True</span><span class="co"> # recommended to reduce vram</span></span>
<span id="cb17-10"><a href="#cb17-10" aria-hidden="true" tabindex="-1"></a></span>
<span id="cb17-11"><a href="#cb17-11" aria-hidden="true" tabindex="-1"></a><span class="fu">datasets</span><span class="kw">:</span></span>
<span id="cb17-12"><a href="#cb17-12" aria-hidden="true" tabindex="-1"></a><span class="at"> </span><span class="kw">-</span><span class="at"> </span><span class="fu">path</span><span class="kw">:</span><span class="at"> ...</span></span>
<span id="cb17-13"><a href="#cb17-13" aria-hidden="true" tabindex="-1"></a><span class="at"> </span><span class="fu">type</span><span class="kw">:</span><span class="at"> </span><span class="st">"axolotl.integrations.kd.chat_template"</span></span>
<span id="cb17-14"><a href="#cb17-14" aria-hidden="true" tabindex="-1"></a><span class="at"> </span><span class="fu">field_messages</span><span class="kw">:</span><span class="at"> </span><span class="st">"messages_combined"</span></span>
<span id="cb17-15"><a href="#cb17-15" aria-hidden="true" tabindex="-1"></a><span class="at"> </span><span class="fu">logprobs_field</span><span class="kw">:</span><span class="at"> </span><span class="st">"llm_text_generation_vllm_logprobs"</span><span class="co"> # for kd only, field of logprobs</span></span></code></pre></div><button title="Copy to Clipboard" class="code-copy-button"><i class="bi"></i></button></div>
<p>An example dataset can be found at <a href="https://huggingface.co/datasets/axolotl-ai-co/evolkit-logprobs-pipeline-75k-v2-sample"><code>axolotl-ai-co/evolkit-logprobs-pipeline-75k-v2-sample</code></a></p>
<p>Please see reference <a href="https://github.com/axolotl-ai-cloud/axolotl/tree/main/src/axolotl/integrations/kd">here</a></p>
</section>
</section>
<section id="llmcompressor" class="level2">
<h2 class="anchored" data-anchor-id="llmcompressor">LLMCompressor</h2>
<p>Fine-tune sparsified models in Axolotl using Neural Magic’s <a href="https://github.com/vllm-project/llm-compressor">LLMCompressor</a>.</p>
<p>This integration enables fine-tuning of models sparsified using LLMCompressor within the Axolotl training framework. By combining LLMCompressor’s model compression capabilities with Axolotl’s distributed training pipelines, users can efficiently fine-tune sparse models at scale.</p>
<p>It uses Axolotl’s plugin system to hook into the fine-tuning flows while maintaining sparsity throughout training.</p>
<hr>
<section id="requirements-1" class="level3">
<h3 class="anchored" data-anchor-id="requirements-1">Requirements</h3>
<ul>
<li><p>Axolotl with <code>llmcompressor</code> extras:</p>
<div class="code-copy-outer-scaffold"><div class="sourceCode" id="cb18"><pre class="sourceCode bash code-with-copy"><code class="sourceCode bash"><span id="cb18-1"><a href="#cb18-1" aria-hidden="true" tabindex="-1"></a><span class="ex">pip</span> install <span class="st">"axolotl[llmcompressor]"</span></span></code></pre></div><button title="Copy to Clipboard" class="code-copy-button"><i class="bi"></i></button></div></li>
<li><p>Requires <code>llmcompressor &gt;= 0.5.1</code></p></li>
</ul>
<p>This will install all necessary dependencies to fine-tune sparsified models using the integration.</p>
<hr>
</section>
<section id="usage-4" class="level3">
<h3 class="anchored" data-anchor-id="usage-4">Usage</h3>
<p>To enable sparse fine-tuning with this integration, include the plugin in your Axolotl config:</p>
<div class="code-copy-outer-scaffold"><div class="sourceCode" id="cb19"><pre class="sourceCode yaml code-with-copy"><code class="sourceCode yaml"><span id="cb19-1"><a href="#cb19-1" aria-hidden="true" tabindex="-1"></a><span class="fu">plugins</span><span class="kw">:</span></span>
<span id="cb19-2"><a href="#cb19-2" aria-hidden="true" tabindex="-1"></a><span class="at"> </span><span class="kw">-</span><span class="at"> axolotl.integrations.llm_compressor.LLMCompressorPlugin</span></span>
<span id="cb19-3"><a href="#cb19-3" aria-hidden="true" tabindex="-1"></a></span>
<span id="cb19-4"><a href="#cb19-4" aria-hidden="true" tabindex="-1"></a><span class="fu">llmcompressor</span><span class="kw">:</span></span>
<span id="cb19-5"><a href="#cb19-5" aria-hidden="true" tabindex="-1"></a><span class="at"> </span><span class="fu">recipe</span><span class="kw">:</span></span>
<span id="cb19-6"><a href="#cb19-6" aria-hidden="true" tabindex="-1"></a><span class="at"> </span><span class="fu">finetuning_stage</span><span class="kw">:</span></span>
<span id="cb19-7"><a href="#cb19-7" aria-hidden="true" tabindex="-1"></a><span class="at"> </span><span class="fu">finetuning_modifiers</span><span class="kw">:</span></span>
<span id="cb19-8"><a href="#cb19-8" aria-hidden="true" tabindex="-1"></a><span class="at"> </span><span class="fu">ConstantPruningModifier</span><span class="kw">:</span></span>
<span id="cb19-9"><a href="#cb19-9" aria-hidden="true" tabindex="-1"></a><span class="at"> </span><span class="fu">targets</span><span class="kw">:</span><span class="at"> </span><span class="kw">[</span></span>
<span id="cb19-10"><a href="#cb19-10" aria-hidden="true" tabindex="-1"></a><span class="at"> </span><span class="st">'re:.*q_proj.weight'</span><span class="kw">,</span></span>
<span id="cb19-11"><a href="#cb19-11" aria-hidden="true" tabindex="-1"></a><span class="at"> </span><span class="st">'re:.*k_proj.weight'</span><span class="kw">,</span></span>
<span id="cb19-12"><a href="#cb19-12" aria-hidden="true" tabindex="-1"></a><span class="at"> </span><span class="st">'re:.*v_proj.weight'</span><span class="kw">,</span></span>
<span id="cb19-13"><a href="#cb19-13" aria-hidden="true" tabindex="-1"></a><span class="at"> </span><span class="st">'re:.*o_proj.weight'</span><span class="kw">,</span></span>
<span id="cb19-14"><a href="#cb19-14" aria-hidden="true" tabindex="-1"></a><span class="at"> </span><span class="st">'re:.*gate_proj.weight'</span><span class="kw">,</span></span>
<span id="cb19-15"><a href="#cb19-15" aria-hidden="true" tabindex="-1"></a><span class="at"> </span><span class="st">'re:.*up_proj.weight'</span><span class="kw">,</span></span>
<span id="cb19-16"><a href="#cb19-16" aria-hidden="true" tabindex="-1"></a><span class="at"> </span><span class="st">'re:.*down_proj.weight'</span><span class="kw">,</span></span>
<span id="cb19-17"><a href="#cb19-17" aria-hidden="true" tabindex="-1"></a><span class="at"> </span><span class="kw">]</span></span>
<span id="cb19-18"><a href="#cb19-18" aria-hidden="true" tabindex="-1"></a><span class="at"> </span><span class="fu">start</span><span class="kw">:</span><span class="at"> </span><span class="dv">0</span></span>
<span id="cb19-19"><a href="#cb19-19" aria-hidden="true" tabindex="-1"></a><span class="at"> </span><span class="fu">save_compressed</span><span class="kw">:</span><span class="at"> </span><span class="ch">true</span></span></code></pre></div><button title="Copy to Clipboard" class="code-copy-button"><i class="bi"></i></button></div>
<p>This plugin <strong>does not apply pruning or sparsification itself</strong> — it is intended for <strong>fine-tuning models that have already been sparsified</strong>.</p>
<p>Pre-sparsified checkpoints can be:
- Generated using <a href="https://github.com/vllm-project/llm-compressor">LLMCompressor</a>
- Downloaded from <a href="https://huggingface.co/neuralmagic">Neural Magic’s Hugging Face page</a>
- Any custom LLM with compatible sparsity patterns that you’ve created yourself</p>
<p>To learn more about writing and customizing LLMCompressor recipes, refer to the official documentation:
<a href="https://github.com/vllm-project/llm-compressor/blob/main/README.md">https://github.com/vllm-project/llm-compressor/blob/main/README.md</a></p>
</section>
<section id="storage-optimization-with-save_compressed" class="level3">
<h3 class="anchored" data-anchor-id="storage-optimization-with-save_compressed">Storage Optimization with save_compressed</h3>
<p>Setting <code>save_compressed: true</code> in your configuration enables saving models in a compressed format, which:
- Reduces disk space usage by approximately 40%
- Maintains compatibility with vLLM for accelerated inference
- Maintains compatibility with llmcompressor for further optimization (example: quantization)</p>
<p>This option is highly recommended when working with sparse models to maximize the benefits of model compression.</p>
</section>
<section id="example-config" class="level3">
<h3 class="anchored" data-anchor-id="example-config">Example Config</h3>
<p>See <a href="examples/llama-3/sparse-finetuning.yaml"><code>examples/llama-3/sparse-finetuning.yaml</code></a> for a complete example.</p>
<hr>
</section>
<section id="inference-with-vllm" class="level3">
<h3 class="anchored" data-anchor-id="inference-with-vllm">Inference with vLLM</h3>
<p>After fine-tuning your sparse model, you can leverage vLLM for efficient inference.
You can also use LLMCompressor to apply additional quantization to your fine-tuned
sparse model before inference for even greater performance benefits.:</p>
<div class="code-copy-outer-scaffold"><div class="sourceCode" id="cb20"><pre class="sourceCode python code-with-copy"><code class="sourceCode python"><span id="cb20-1"><a href="#cb20-1" aria-hidden="true" tabindex="-1"></a><span class="im">from</span> vllm <span class="im">import</span> LLM, SamplingParams</span>
<span id="cb20-2"><a href="#cb20-2" aria-hidden="true" tabindex="-1"></a></span>
<span id="cb20-3"><a href="#cb20-3" aria-hidden="true" tabindex="-1"></a>prompts <span class="op">=</span> [</span>
<span id="cb20-4"><a href="#cb20-4" aria-hidden="true" tabindex="-1"></a> <span class="st">"Hello, my name is"</span>,</span>
<span id="cb20-5"><a href="#cb20-5" aria-hidden="true" tabindex="-1"></a> <span class="st">"The president of the United States is"</span>,</span>
<span id="cb20-6"><a href="#cb20-6" aria-hidden="true" tabindex="-1"></a> <span class="st">"The capital of France is"</span>,</span>
<span id="cb20-7"><a href="#cb20-7" aria-hidden="true" tabindex="-1"></a> <span class="st">"The future of AI is"</span>,</span>
<span id="cb20-8"><a href="#cb20-8" aria-hidden="true" tabindex="-1"></a>]</span>
<span id="cb20-9"><a href="#cb20-9" aria-hidden="true" tabindex="-1"></a>sampling_params <span class="op">=</span> SamplingParams(temperature<span class="op">=</span><span class="fl">0.8</span>, top_p<span class="op">=</span><span class="fl">0.95</span>)</span>
<span id="cb20-10"><a href="#cb20-10" aria-hidden="true" tabindex="-1"></a>llm <span class="op">=</span> LLM(<span class="st">"path/to/your/sparse/model"</span>)</span>
<span id="cb20-11"><a href="#cb20-11" aria-hidden="true" tabindex="-1"></a>outputs <span class="op">=</span> llm.generate(prompts, sampling_params)</span>
<span id="cb20-12"><a href="#cb20-12" aria-hidden="true" tabindex="-1"></a></span>
<span id="cb20-13"><a href="#cb20-13" aria-hidden="true" tabindex="-1"></a><span class="cf">for</span> output <span class="kw">in</span> outputs:</span>
<span id="cb20-14"><a href="#cb20-14" aria-hidden="true" tabindex="-1"></a> prompt <span class="op">=</span> output.prompt</span>
<span id="cb20-15"><a href="#cb20-15" aria-hidden="true" tabindex="-1"></a> generated_text <span class="op">=</span> output.outputs[<span class="dv">0</span>].text</span>
<span id="cb20-16"><a href="#cb20-16" aria-hidden="true" tabindex="-1"></a> <span class="bu">print</span>(<span class="ss">f"Prompt: </span><span class="sc">{</span>prompt<span class="sc">!r}</span><span class="ss">, Generated text: </span><span class="sc">{</span>generated_text<span class="sc">!r}</span><span class="ss">"</span>)</span></code></pre></div><button title="Copy to Clipboard" class="code-copy-button"><i class="bi"></i></button></div>
<p>For more details on vLLM’s capabilities and advanced configuration options, see the <a href="https://docs.vllm.ai/">official vLLM documentation</a>.</p>
</section>
<section id="learn-more" class="level3">
<h3 class="anchored" data-anchor-id="learn-more">Learn More</h3>
<p>For details on available sparsity and quantization schemes, fine-tuning recipes, and usage examples, visit the official LLMCompressor repository:</p>
<p><a href="https://github.com/vllm-project/llm-compressor">https://github.com/vllm-project/llm-compressor</a></p>
<p>Please see reference <a href="https://github.com/axolotl-ai-cloud/axolotl/tree/main/src/axolotl/integrations/llm_compressor">here</a></p>
</section>
</section>
<section id="language-model-evaluation-harness-lm-eval" class="level2">
<h2 class="anchored" data-anchor-id="language-model-evaluation-harness-lm-eval">Language Model Evaluation Harness (LM Eval)</h2>
<p>Run evaluation on model using the popular lm-evaluation-harness library.</p>
<p>See https://github.com/EleutherAI/lm-evaluation-harness</p>
<section id="usage-5" class="level3">
<h3 class="anchored" data-anchor-id="usage-5">Usage</h3>
<p>There are two ways to use the LM Eval integration:</p>
</section>
<section id="post-training-evaluation" class="level3">
<h3 class="anchored" data-anchor-id="post-training-evaluation">1. Post-Training Evaluation</h3>
<p>When training with the plugin enabled, evaluation runs automatically after training completes:</p>
<div class="code-copy-outer-scaffold"><div class="sourceCode" id="cb21"><pre class="sourceCode yaml code-with-copy"><code class="sourceCode yaml"><span id="cb21-1"><a href="#cb21-1" aria-hidden="true" tabindex="-1"></a><span class="fu">plugins</span><span class="kw">:</span></span>
<span id="cb21-2"><a href="#cb21-2" aria-hidden="true" tabindex="-1"></a><span class="at"> </span><span class="kw">-</span><span class="at"> axolotl.integrations.lm_eval.LMEvalPlugin</span></span>
<span id="cb21-3"><a href="#cb21-3" aria-hidden="true" tabindex="-1"></a></span>
<span id="cb21-4"><a href="#cb21-4" aria-hidden="true" tabindex="-1"></a><span class="fu">lm_eval_tasks</span><span class="kw">:</span></span>
<span id="cb21-5"><a href="#cb21-5" aria-hidden="true" tabindex="-1"></a><span class="at"> </span><span class="kw">-</span><span class="at"> gsm8k</span></span>
<span id="cb21-6"><a href="#cb21-6" aria-hidden="true" tabindex="-1"></a><span class="at"> </span><span class="kw">-</span><span class="at"> hellaswag</span></span>
<span id="cb21-7"><a href="#cb21-7" aria-hidden="true" tabindex="-1"></a><span class="at"> </span><span class="kw">-</span><span class="at"> arc_easy</span></span>
<span id="cb21-8"><a href="#cb21-8" aria-hidden="true" tabindex="-1"></a></span>
<span id="cb21-9"><a href="#cb21-9" aria-hidden="true" tabindex="-1"></a><span class="fu">lm_eval_batch_size</span><span class="kw">:</span><span class="co"> # Batch size for evaluation</span></span>
<span id="cb21-10"><a href="#cb21-10" aria-hidden="true" tabindex="-1"></a></span>
<span id="cb21-11"><a href="#cb21-11" aria-hidden="true" tabindex="-1"></a><span class="fu">output_dir</span><span class="kw">:</span></span></code></pre></div><button title="Copy to Clipboard" class="code-copy-button"><i class="bi"></i></button></div>
<p>Run training as usual:</p>
<div class="code-copy-outer-scaffold"><div class="sourceCode" id="cb22"><pre class="sourceCode bash code-with-copy"><code class="sourceCode bash"><span id="cb22-1"><a href="#cb22-1" aria-hidden="true" tabindex="-1"></a><span class="ex">axolotl</span> train config.yml</span></code></pre></div><button title="Copy to Clipboard" class="code-copy-button"><i class="bi"></i></button></div>
</section>
<section id="standalone-cli-evaluation" class="level3">
<h3 class="anchored" data-anchor-id="standalone-cli-evaluation">2. Standalone CLI Evaluation</h3>
<p>Evaluate any model directly without training:</p>
<div class="code-copy-outer-scaffold"><div class="sourceCode" id="cb23"><pre class="sourceCode yaml code-with-copy"><code class="sourceCode yaml"><span id="cb23-1"><a href="#cb23-1" aria-hidden="true" tabindex="-1"></a><span class="fu">lm_eval_model</span><span class="kw">:</span><span class="at"> meta-llama/Llama-2-7b-hf</span></span>
<span id="cb23-2"><a href="#cb23-2" aria-hidden="true" tabindex="-1"></a></span>
<span id="cb23-3"><a href="#cb23-3" aria-hidden="true" tabindex="-1"></a><span class="fu">plugins</span><span class="kw">:</span></span>
<span id="cb23-4"><a href="#cb23-4" aria-hidden="true" tabindex="-1"></a><span class="at"> </span><span class="kw">-</span><span class="at"> axolotl.integrations.lm_eval.LMEvalPlugin</span></span>
<span id="cb23-5"><a href="#cb23-5" aria-hidden="true" tabindex="-1"></a></span>
<span id="cb23-6"><a href="#cb23-6" aria-hidden="true" tabindex="-1"></a><span class="fu">lm_eval_tasks</span><span class="kw">:</span></span>
<span id="cb23-7"><a href="#cb23-7" aria-hidden="true" tabindex="-1"></a><span class="at"> </span><span class="kw">-</span><span class="at"> gsm8k</span></span>
<span id="cb23-8"><a href="#cb23-8" aria-hidden="true" tabindex="-1"></a><span class="at"> </span><span class="kw">-</span><span class="at"> hellaswag</span></span>
<span id="cb23-9"><a href="#cb23-9" aria-hidden="true" tabindex="-1"></a><span class="at"> </span><span class="kw">-</span><span class="at"> arc_easy</span></span>
<span id="cb23-10"><a href="#cb23-10" aria-hidden="true" tabindex="-1"></a></span>
<span id="cb23-11"><a href="#cb23-11" aria-hidden="true" tabindex="-1"></a><span class="fu">lm_eval_batch_size</span><span class="kw">:</span><span class="at"> </span><span class="dv">8</span></span>
<span id="cb23-12"><a href="#cb23-12" aria-hidden="true" tabindex="-1"></a><span class="fu">output_dir</span><span class="kw">:</span><span class="at"> ./outputs</span></span></code></pre></div><button title="Copy to Clipboard" class="code-copy-button"><i class="bi"></i></button></div>
<p>Run evaluation:</p>
<div class="code-copy-outer-scaffold"><div class="sourceCode" id="cb24"><pre class="sourceCode bash code-with-copy"><code class="sourceCode bash"><span id="cb24-1"><a href="#cb24-1" aria-hidden="true" tabindex="-1"></a><span class="ex">axolotl</span> lm-eval config.yml</span></code></pre></div><button title="Copy to Clipboard" class="code-copy-button"><i class="bi"></i></button></div>
</section>
<section id="model-selection-priority" class="level3">
<h3 class="anchored" data-anchor-id="model-selection-priority">Model Selection Priority</h3>
<p>The model to evaluate is selected in the following priority order:</p>
<ol type="1">
<li><strong><code>lm_eval_model</code></strong> - Explicit model path or HuggingFace repo (highest priority)</li>
<li><strong><code>hub_model_id</code></strong> - Trained model pushed to HuggingFace Hub</li>
<li><strong><code>output_dir</code></strong> - Local checkpoint directory containing trained model weights</li>
</ol>
</section>
<section id="citation-2" class="level3">
<h3 class="anchored" data-anchor-id="citation-2">Citation</h3>
<div class="code-copy-outer-scaffold"><div class="sourceCode" id="cb25"><pre class="sourceCode bib code-with-copy"><code class="sourceCode bibtex"><span id="cb25-1"><a href="#cb25-1" aria-hidden="true" tabindex="-1"></a><span class="va">@misc</span>{<span class="ot">eval</span>-<span class="ot">harness</span>,</span>
<span id="cb25-2"><a href="#cb25-2" aria-hidden="true" tabindex="-1"></a> <span class="dt">author</span> = {Gao, Leo and Tow, Jonathan and Abbasi, Baber and Biderman, Stella and Black, Sid and DiPofi, Anthony and Foster, Charles and Golding, Laurence and Hsu, Jeffrey and Le Noac'h, Alain and Li, Haonan and McDonell, Kyle and Muennighoff, Niklas and Ociepa, Chris and Phang, Jason and Reynolds, Laria and Schoelkopf, Hailey and Skowron, Aviya and Sutawika, Lintang and Tang, Eric and Thite, Anish and Wang, Ben and Wang, Kevin and Zou, Andy},</span>
<span id="cb25-3"><a href="#cb25-3" aria-hidden="true" tabindex="-1"></a> <span class="dt">title</span> = {A framework for few-shot language model evaluation},</span>
<span id="cb25-4"><a href="#cb25-4" aria-hidden="true" tabindex="-1"></a> <span class="dt">month</span> = 07,</span>
<span id="cb25-5"><a href="#cb25-5" aria-hidden="true" tabindex="-1"></a> <span class="dt">year</span> = 2024,</span>
<span id="cb25-6"><a href="#cb25-6" aria-hidden="true" tabindex="-1"></a> <span class="dt">publisher</span> = {Zenodo},</span>
<span id="cb25-7"><a href="#cb25-7" aria-hidden="true" tabindex="-1"></a> <span class="dt">version</span> = {v0.4.3},</span>
<span id="cb25-8"><a href="#cb25-8" aria-hidden="true" tabindex="-1"></a> <span class="dt">doi</span> = {10.5281/zenodo.12608602},</span>
<span id="cb25-9"><a href="#cb25-9" aria-hidden="true" tabindex="-1"></a> <span class="dt">url</span> = {https://zenodo.org/records/12608602}</span>
<span id="cb25-10"><a href="#cb25-10" aria-hidden="true" tabindex="-1"></a>}</span></code></pre></div><button title="Copy to Clipboard" class="code-copy-button"><i class="bi"></i></button></div>
<p>Please see reference <a href="https://github.com/axolotl-ai-cloud/axolotl/tree/main/src/axolotl/integrations/lm_eval">here</a></p>
</section>
</section>
<section id="liger-kernels" class="level2">
<h2 class="anchored" data-anchor-id="liger-kernels">Liger Kernels</h2>
<p>Liger Kernel provides efficient Triton kernels for LLM training, offering:</p>
<ul>
<li>20% increase in multi-GPU training throughput</li>
<li>60% reduction in memory usage</li>
<li>Compatibility with both FSDP and DeepSpeed</li>
</ul>
<p>See https://github.com/linkedin/Liger-Kernel</p>
<section id="usage-6" class="level3">
<h3 class="anchored" data-anchor-id="usage-6">Usage</h3>
<div class="code-copy-outer-scaffold"><div class="sourceCode" id="cb26"><pre class="sourceCode yaml code-with-copy"><code class="sourceCode yaml"><span id="cb26-1"><a href="#cb26-1" aria-hidden="true" tabindex="-1"></a><span class="fu">plugins</span><span class="kw">:</span></span>
<span id="cb26-2"><a href="#cb26-2" aria-hidden="true" tabindex="-1"></a><span class="at"> </span><span class="kw">-</span><span class="at"> axolotl.integrations.liger.LigerPlugin</span></span>
<span id="cb26-3"><a href="#cb26-3" aria-hidden="true" tabindex="-1"></a><span class="fu">liger_rope</span><span class="kw">:</span><span class="at"> </span><span class="ch">true</span></span>
<span id="cb26-4"><a href="#cb26-4" aria-hidden="true" tabindex="-1"></a><span class="fu">liger_rms_norm</span><span class="kw">:</span><span class="at"> </span><span class="ch">true</span></span>
<span id="cb26-5"><a href="#cb26-5" aria-hidden="true" tabindex="-1"></a><span class="fu">liger_glu_activation</span><span class="kw">:</span><span class="at"> </span><span class="ch">true</span></span>
<span id="cb26-6"><a href="#cb26-6" aria-hidden="true" tabindex="-1"></a><span class="fu">liger_layer_norm</span><span class="kw">:</span><span class="at"> </span><span class="ch">true</span></span>
<span id="cb26-7"><a href="#cb26-7" aria-hidden="true" tabindex="-1"></a><span class="fu">liger_fused_linear_cross_entropy</span><span class="kw">:</span><span class="at"> </span><span class="ch">true</span></span>
<span id="cb26-8"><a href="#cb26-8" aria-hidden="true" tabindex="-1"></a></span>
<span id="cb26-9"><a href="#cb26-9" aria-hidden="true" tabindex="-1"></a><span class="fu">liger_use_token_scaling</span><span class="kw">:</span><span class="at"> </span><span class="ch">true</span></span></code></pre></div><button title="Copy to Clipboard" class="code-copy-button"><i class="bi"></i></button></div>
</section>
<section id="supported-models-2" class="level3">
<h3 class="anchored" data-anchor-id="supported-models-2">Supported Models</h3>
<ul>
<li>deepseek_v2</li>
<li>gemma</li>
<li>gemma2</li>
<li>gemma3</li>
<li>granite</li>
<li>jamba</li>
<li>llama</li>
<li>mistral</li>
<li>mixtral</li>
<li>mllama</li>
<li>mllama_text_model</li>
<li>olmo2</li>
<li>paligemma</li>
<li>phi3</li>
<li>qwen2</li>
<li>qwen2_5_vl</li>
<li>qwen2_vl</li>
</ul>
</section>
<section id="citation-3" class="level3">
<h3 class="anchored" data-anchor-id="citation-3">Citation</h3>
<div class="code-copy-outer-scaffold"><div class="sourceCode" id="cb27"><pre class="sourceCode bib code-with-copy"><code class="sourceCode bibtex"><span id="cb27-1"><a href="#cb27-1" aria-hidden="true" tabindex="-1"></a><span class="va">@article</span>{<span class="ot">hsu2024ligerkernelefficienttriton</span>,</span>
<span id="cb27-2"><a href="#cb27-2" aria-hidden="true" tabindex="-1"></a> <span class="dt">title</span>={Liger Kernel: Efficient Triton Kernels for LLM Training},</span>
<span id="cb27-3"><a href="#cb27-3" aria-hidden="true" tabindex="-1"></a> <span class="dt">author</span>={Pin-Lun Hsu and Yun Dai and Vignesh Kothapalli and Qingquan Song and Shao Tang and Siyu Zhu and Steven Shimizu and Shivam Sahni and Haowen Ning and Yanning Chen},</span>
<span id="cb27-4"><a href="#cb27-4" aria-hidden="true" tabindex="-1"></a> <span class="dt">year</span>={2024},</span>
<span id="cb27-5"><a href="#cb27-5" aria-hidden="true" tabindex="-1"></a> <span class="dt">eprint</span>={2410.10989},</span>
<span id="cb27-6"><a href="#cb27-6" aria-hidden="true" tabindex="-1"></a> <span class="dt">archivePrefix</span>={arXiv},</span>
<span id="cb27-7"><a href="#cb27-7" aria-hidden="true" tabindex="-1"></a> <span class="dt">primaryClass</span>={cs.LG},</span>
<span id="cb27-8"><a href="#cb27-8" aria-hidden="true" tabindex="-1"></a> <span class="dt">url</span>={https://arxiv.org/abs/2410.10989},</span>
<span id="cb27-9"><a href="#cb27-9" aria-hidden="true" tabindex="-1"></a> <span class="dt">journal</span>={arXiv preprint arXiv:2410.10989},</span>
<span id="cb27-10"><a href="#cb27-10" aria-hidden="true" tabindex="-1"></a>}</span></code></pre></div><button title="Copy to Clipboard" class="code-copy-button"><i class="bi"></i></button></div>
<p>Please see reference <a href="https://github.com/axolotl-ai-cloud/axolotl/tree/main/src/axolotl/integrations/liger">here</a></p>
</section>
</section>
<section id="nemo-gym-integration-for-axolotl" class="level2">
<h2 class="anchored" data-anchor-id="nemo-gym-integration-for-axolotl">NeMo Gym Integration for Axolotl</h2>
<p>Train LLMs with reinforcement learning using <a href="https://github.com/NVIDIA-NeMo/Gym">NVIDIA NeMo Gym</a> environments as reward sources. NeMo Gym provides 50+ verified RL environments spanning math, coding, tool-use, reasoning, and safety — each with deterministic reward signals.</p>
<section id="validated-training-paths" class="level3">
<h3 class="anchored" data-anchor-id="validated-training-paths">Validated Training Paths</h3>
<table class="caption-top table">
<colgroup>
<col style="width: 15%">
<col style="width: 17%">
<col style="width: 30%">
<col style="width: 35%">
</colgroup>
<thead>
<tr class="header">
<th>Path</th>
<th>Speed</th>
<th>Multi-turn</th>
<th>Architecture</th>
</tr>
</thead>
<tbody>
<tr class="odd">
<td><strong>Async GRPO + Data Producer</strong></td>
<td>Fastest (3x)</td>
<td>Yes</td>
<td><code>NemoGymDataProducer</code> replaces vLLM generation</td>
</tr>
<tr class="even">
<td>Standard GRPO + Data Producer</td>
<td>Baseline</td>
<td>Yes</td>
<td>Same producer, no async prefetch</td>
</tr>
<tr class="odd">
<td>Standard GRPO + /verify</td>
<td>Simplest</td>
<td>No</td>
<td>Reward function calls /verify directly</td>
</tr>
<tr class="even">
<td>FSDP2 + /verify (2 GPU)</td>
<td>Distributed</td>
<td>No</td>
<td><code>fsdp_version: 2</code></td>
</tr>
</tbody>
</table>
<p>Multi-turn uses <code>nemo_gym_multi_turn: true</code> which auto-enables the async trainer’s
data producer protocol. The plugin’s <code>NemoGymDataProducer</code> calls NeMo Gym agent <code>/run</code>
endpoints and returns <code>RolloutDataset</code> with proper IS correction, env_mask, and rewards.</p>
<p>All paths tested end-to-end with Qwen3-0.6B + LoRA, logged to wandb project <code>nemo-gym-rl</code>.</p>
</section>
<section id="quick-start" class="level3">
<h3 class="anchored" data-anchor-id="quick-start">Quick Start</h3>
</section>
<section id="prerequisites" class="level3">
<h3 class="anchored" data-anchor-id="prerequisites">Prerequisites</h3>
<ul>
<li><a href="https://github.com/astral-sh/uv">uv</a> package manager (for NeMo Gym’s venv)</li>
<li>Two GPUs recommended (one for vLLM server, one for training)</li>
</ul>
</section>
<section id="set-up-nemo-gym" class="level3">
<h3 class="anchored" data-anchor-id="set-up-nemo-gym">1. Set Up NeMo Gym</h3>
<div class="code-copy-outer-scaffold"><div class="sourceCode" id="cb28"><pre class="sourceCode bash code-with-copy"><code class="sourceCode bash"><span id="cb28-1"><a href="#cb28-1" aria-hidden="true" tabindex="-1"></a><span class="fu">git</span> clone https://github.com/NVIDIA-NeMo/Gym.git ~/Gym</span>
<span id="cb28-2"><a href="#cb28-2" aria-hidden="true" tabindex="-1"></a><span class="bu">cd</span> ~/Gym</span>
<span id="cb28-3"><a href="#cb28-3" aria-hidden="true" tabindex="-1"></a><span class="ex">uv</span> venv <span class="at">--python</span> 3.12 <span class="kw">&amp;&amp;</span> <span class="bu">source</span> .venv/bin/activate <span class="kw">&amp;&amp;</span> <span class="ex">uv</span> sync</span>
<span id="cb28-4"><a href="#cb28-4" aria-hidden="true" tabindex="-1"></a></span>
<span id="cb28-5"><a href="#cb28-5" aria-hidden="true" tabindex="-1"></a><span class="va">CFLAGS</span><span class="op">=</span><span class="st">""</span> <span class="ex">uv</span> pip install pycosat <span class="at">--python</span> .venv/bin/python <span class="at">--no-build-isolation</span></span>
<span id="cb28-6"><a href="#cb28-6" aria-hidden="true" tabindex="-1"></a></span>
<span id="cb28-7"><a href="#cb28-7" aria-hidden="true" tabindex="-1"></a><span class="cf">for</span> dir <span class="kw">in</span> resources_servers/reasoning_gym resources_servers/example_single_tool_call responses_api_models/vllm_model responses_api_agents/simple_agent<span class="kw">;</span> <span class="cf">do</span></span>
<span id="cb28-8"><a href="#cb28-8" aria-hidden="true" tabindex="-1"></a> <span class="ex">uv</span> venv <span class="at">--seed</span> <span class="at">--allow-existing</span> <span class="at">--python</span> 3.12 <span class="va">$dir</span>/.venv</span>
<span id="cb28-9"><a href="#cb28-9" aria-hidden="true" tabindex="-1"></a> <span class="va">CFLAGS</span><span class="op">=</span><span class="st">""</span> <span class="ex">uv</span> pip install <span class="at">--python</span> <span class="va">$dir</span>/.venv/bin/python pycosat <span class="at">--no-build-isolation</span> <span class="dv">2</span><span class="op">&gt;</span>/dev/null</span>
<span id="cb28-10"><a href="#cb28-10" aria-hidden="true" tabindex="-1"></a> <span class="ex">uv</span> pip install <span class="at">--python</span> <span class="va">$dir</span>/.venv/bin/python <span class="at">-e</span> . <span class="st">"ray[default]==2.52.1"</span></span>
<span id="cb28-11"><a href="#cb28-11" aria-hidden="true" tabindex="-1"></a><span class="cf">done</span></span>
<span id="cb28-12"><a href="#cb28-12" aria-hidden="true" tabindex="-1"></a></span>
<span id="cb28-13"><a href="#cb28-13" aria-hidden="true" tabindex="-1"></a><span class="ex">uv</span> pip install <span class="at">--python</span> resources_servers/reasoning_gym/.venv/bin/python <span class="dt">\</span></span>
<span id="cb28-14"><a href="#cb28-14" aria-hidden="true" tabindex="-1"></a> reasoning-gym matplotlib pillow cycler contourpy kiwisolver</span></code></pre></div><button title="Copy to Clipboard" class="code-copy-button"><i class="bi"></i></button></div>
</section>
<section id="multi-turn-with-async-grpo-recommended-fastest-path" class="level3">
<h3 class="anchored" data-anchor-id="multi-turn-with-async-grpo-recommended-fastest-path">2. Multi-Turn with Async GRPO (Recommended — Fastest Path)</h3>
<p>This is the fully validated, highest-performance path. NeMo Gym’s agent server handles
multi-turn tool execution while axolotl’s async GRPO prefetches data in background threads.</p>
<p><strong>Step 1: Create the NeMo Gym agent config</strong></p>
<p>Create <code>~/Gym/configs/axolotl_tool_calling.yaml</code>:</p>
<div class="code-copy-outer-scaffold"><div class="sourceCode" id="cb29"><pre class="sourceCode yaml code-with-copy"><code class="sourceCode yaml"><span id="cb29-1"><a href="#cb29-1" aria-hidden="true" tabindex="-1"></a><span class="fu">example_single_tool_call</span><span class="kw">:</span></span>
<span id="cb29-2"><a href="#cb29-2" aria-hidden="true" tabindex="-1"></a><span class="at"> </span><span class="fu">resources_servers</span><span class="kw">:</span></span>
<span id="cb29-3"><a href="#cb29-3" aria-hidden="true" tabindex="-1"></a><span class="at"> </span><span class="fu">example_single_tool_call</span><span class="kw">:</span></span>
<span id="cb29-4"><a href="#cb29-4" aria-hidden="true" tabindex="-1"></a><span class="at"> </span><span class="fu">entrypoint</span><span class="kw">:</span><span class="at"> app.py</span></span>
<span id="cb29-5"><a href="#cb29-5" aria-hidden="true" tabindex="-1"></a><span class="at"> </span><span class="fu">domain</span><span class="kw">:</span><span class="at"> agent</span></span>
<span id="cb29-6"><a href="#cb29-6" aria-hidden="true" tabindex="-1"></a><span class="at"> </span><span class="fu">verified</span><span class="kw">:</span><span class="at"> </span><span class="ch">false</span></span>
<span id="cb29-7"><a href="#cb29-7" aria-hidden="true" tabindex="-1"></a></span>
<span id="cb29-8"><a href="#cb29-8" aria-hidden="true" tabindex="-1"></a><span class="fu">policy_model</span><span class="kw">:</span></span>
<span id="cb29-9"><a href="#cb29-9" aria-hidden="true" tabindex="-1"></a><span class="at"> </span><span class="fu">responses_api_models</span><span class="kw">:</span></span>
<span id="cb29-10"><a href="#cb29-10" aria-hidden="true" tabindex="-1"></a><span class="at"> </span><span class="fu">vllm_model</span><span class="kw">:</span></span>
<span id="cb29-11"><a href="#cb29-11" aria-hidden="true" tabindex="-1"></a><span class="at"> </span><span class="fu">entrypoint</span><span class="kw">:</span><span class="at"> app.py</span></span>
<span id="cb29-12"><a href="#cb29-12" aria-hidden="true" tabindex="-1"></a><span class="at"> </span><span class="fu">base_url</span><span class="kw">:</span><span class="at"> http://localhost:8000/v1</span></span>
<span id="cb29-13"><a href="#cb29-13" aria-hidden="true" tabindex="-1"></a><span class="at"> </span><span class="fu">api_key</span><span class="kw">:</span><span class="at"> dummy_key</span></span>
<span id="cb29-14"><a href="#cb29-14" aria-hidden="true" tabindex="-1"></a><span class="at"> </span><span class="fu">model</span><span class="kw">:</span><span class="at"> Qwen/Qwen3-0.6B</span><span class="co"> # Must match your training model</span></span>
<span id="cb29-15"><a href="#cb29-15" aria-hidden="true" tabindex="-1"></a><span class="at"> </span><span class="fu">return_token_id_information</span><span class="kw">:</span><span class="at"> </span><span class="ch">true</span></span>
<span id="cb29-16"><a href="#cb29-16" aria-hidden="true" tabindex="-1"></a><span class="at"> </span><span class="fu">uses_reasoning_parser</span><span class="kw">:</span><span class="at"> </span><span class="ch">false</span></span>
<span id="cb29-17"><a href="#cb29-17" aria-hidden="true" tabindex="-1"></a></span>
<span id="cb29-18"><a href="#cb29-18" aria-hidden="true" tabindex="-1"></a><span class="fu">example_single_tool_call_simple_agent</span><span class="kw">:</span></span>
<span id="cb29-19"><a href="#cb29-19" aria-hidden="true" tabindex="-1"></a><span class="at"> </span><span class="fu">responses_api_agents</span><span class="kw">:</span></span>
<span id="cb29-20"><a href="#cb29-20" aria-hidden="true" tabindex="-1"></a><span class="at"> </span><span class="fu">simple_agent</span><span class="kw">:</span></span>
<span id="cb29-21"><a href="#cb29-21" aria-hidden="true" tabindex="-1"></a><span class="at"> </span><span class="fu">entrypoint</span><span class="kw">:</span><span class="at"> app.py</span></span>
<span id="cb29-22"><a href="#cb29-22" aria-hidden="true" tabindex="-1"></a><span class="at"> </span><span class="fu">resources_server</span><span class="kw">:</span></span>
<span id="cb29-23"><a href="#cb29-23" aria-hidden="true" tabindex="-1"></a><span class="at"> </span><span class="fu">type</span><span class="kw">:</span><span class="at"> resources_servers</span></span>
<span id="cb29-24"><a href="#cb29-24" aria-hidden="true" tabindex="-1"></a><span class="at"> </span><span class="fu">name</span><span class="kw">:</span><span class="at"> example_single_tool_call</span></span>
<span id="cb29-25"><a href="#cb29-25" aria-hidden="true" tabindex="-1"></a><span class="at"> </span><span class="fu">model_server</span><span class="kw">:</span></span>
<span id="cb29-26"><a href="#cb29-26" aria-hidden="true" tabindex="-1"></a><span class="at"> </span><span class="fu">type</span><span class="kw">:</span><span class="at"> responses_api_models</span></span>
<span id="cb29-27"><a href="#cb29-27" aria-hidden="true" tabindex="-1"></a><span class="at"> </span><span class="fu">name</span><span class="kw">:</span><span class="at"> policy_model</span></span>
<span id="cb29-28"><a href="#cb29-28" aria-hidden="true" tabindex="-1"></a><span class="at"> </span><span class="fu">datasets</span><span class="kw">:</span></span>
<span id="cb29-29"><a href="#cb29-29" aria-hidden="true" tabindex="-1"></a><span class="at"> </span><span class="kw">-</span><span class="at"> </span><span class="fu">name</span><span class="kw">:</span><span class="at"> weather</span></span>
<span id="cb29-30"><a href="#cb29-30" aria-hidden="true" tabindex="-1"></a><span class="at"> </span><span class="fu">type</span><span class="kw">:</span><span class="at"> example</span></span>
<span id="cb29-31"><a href="#cb29-31" aria-hidden="true" tabindex="-1"></a><span class="at"> </span><span class="fu">jsonl_fpath</span><span class="kw">:</span><span class="at"> resources_servers/example_single_tool_call/data/weather_tool_calling.jsonl</span></span></code></pre></div><button title="Copy to Clipboard" class="code-copy-button"><i class="bi"></i></button></div>
<p><strong>Step 2: Start three services</strong></p>
<div class="code-copy-outer-scaffold"><div class="sourceCode" id="cb30"><pre class="sourceCode bash code-with-copy"><code class="sourceCode bash"><span id="cb30-1"><a href="#cb30-1" aria-hidden="true" tabindex="-1"></a><span class="va">CUDA_VISIBLE_DEVICES</span><span class="op">=</span>0 <span class="ex">python</span> <span class="at">-m</span> vllm.entrypoints.openai.api_server <span class="dt">\</span></span>
<span id="cb30-2"><a href="#cb30-2" aria-hidden="true" tabindex="-1"></a> <span class="at">--model</span> Qwen/Qwen3-0.6B <span class="at">--max-model-len</span> 2048 <span class="at">--gpu-memory-utilization</span> 0.85</span>
<span id="cb30-3"><a href="#cb30-3" aria-hidden="true" tabindex="-1"></a></span>
<span id="cb30-4"><a href="#cb30-4" aria-hidden="true" tabindex="-1"></a><span class="bu">cd</span> ~/Gym <span class="kw">&amp;&amp;</span> <span class="ex">.venv/bin/ng_run</span> <span class="dt">\</span></span>
<span id="cb30-5"><a href="#cb30-5" aria-hidden="true" tabindex="-1"></a> <span class="st">"+config_paths=[configs/axolotl_tool_calling.yaml]"</span> <span class="st">"+skip_venv_if_present=true"</span></span>
<span id="cb30-6"><a href="#cb30-6" aria-hidden="true" tabindex="-1"></a></span>
<span id="cb30-7"><a href="#cb30-7" aria-hidden="true" tabindex="-1"></a><span class="bu">cd</span> experiments <span class="kw">&amp;&amp;</span> <span class="va">CUDA_VISIBLE_DEVICES</span><span class="op">=</span>1 <span class="va">CUDA_HOME</span><span class="op">=</span><span class="va">$HOME</span>/env-claude-cu130/cuda_shim <span class="dt">\</span></span>
<span id="cb30-8"><a href="#cb30-8" aria-hidden="true" tabindex="-1"></a> <span class="ex">axolotl</span> train nemo_gym_async_agent.yaml</span></code></pre></div><button title="Copy to Clipboard" class="code-copy-button"><i class="bi"></i></button></div>
<p><strong>Step 3: Training config</strong> (<code>nemo_gym_async_agent.yaml</code>):</p>
<div class="code-copy-outer-scaffold"><div class="sourceCode" id="cb31"><pre class="sourceCode yaml code-with-copy"><code class="sourceCode yaml"><span id="cb31-1"><a href="#cb31-1" aria-hidden="true" tabindex="-1"></a><span class="fu">base_model</span><span class="kw">:</span><span class="at"> Qwen/Qwen3-0.6B</span></span>
<span id="cb31-2"><a href="#cb31-2" aria-hidden="true" tabindex="-1"></a><span class="fu">adapter</span><span class="kw">:</span><span class="at"> lora</span></span>
<span id="cb31-3"><a href="#cb31-3" aria-hidden="true" tabindex="-1"></a><span class="fu">lora_r</span><span class="kw">:</span><span class="at"> </span><span class="dv">16</span></span>
<span id="cb31-4"><a href="#cb31-4" aria-hidden="true" tabindex="-1"></a><span class="fu">lora_alpha</span><span class="kw">:</span><span class="at"> </span><span class="dv">32</span></span>
<span id="cb31-5"><a href="#cb31-5" aria-hidden="true" tabindex="-1"></a><span class="fu">lora_target_modules</span><span class="kw">:</span><span class="at"> </span><span class="kw">[</span><span class="at">q_proj</span><span class="kw">,</span><span class="at"> k_proj</span><span class="kw">,</span><span class="at"> v_proj</span><span class="kw">,</span><span class="at"> o_proj</span><span class="kw">,</span><span class="at"> gate_proj</span><span class="kw">,</span><span class="at"> up_proj</span><span class="kw">,</span><span class="at"> down_proj</span><span class="kw">]</span></span>
<span id="cb31-6"><a href="#cb31-6" aria-hidden="true" tabindex="-1"></a><span class="fu">sequence_len</span><span class="kw">:</span><span class="at"> </span><span class="dv">2048</span></span>
<span id="cb31-7"><a href="#cb31-7" aria-hidden="true" tabindex="-1"></a></span>
<span id="cb31-8"><a href="#cb31-8" aria-hidden="true" tabindex="-1"></a><span class="fu">rl</span><span class="kw">:</span><span class="at"> grpo</span></span>
<span id="cb31-9"><a href="#cb31-9" aria-hidden="true" tabindex="-1"></a><span class="fu">chat_template</span><span class="kw">:</span><span class="at"> tokenizer_default</span></span>
<span id="cb31-10"><a href="#cb31-10" aria-hidden="true" tabindex="-1"></a></span>
<span id="cb31-11"><a href="#cb31-11" aria-hidden="true" tabindex="-1"></a><span class="fu">trl</span><span class="kw">:</span></span>
<span id="cb31-12"><a href="#cb31-12" aria-hidden="true" tabindex="-1"></a><span class="at"> </span><span class="fu">use_vllm</span><span class="kw">:</span><span class="at"> </span><span class="ch">true</span></span>
<span id="cb31-13"><a href="#cb31-13" aria-hidden="true" tabindex="-1"></a><span class="at"> </span><span class="fu">vllm_mode</span><span class="kw">:</span><span class="at"> server</span></span>
<span id="cb31-14"><a href="#cb31-14" aria-hidden="true" tabindex="-1"></a><span class="at"> </span><span class="fu">vllm_server_host</span><span class="kw">:</span><span class="at"> localhost</span></span>
<span id="cb31-15"><a href="#cb31-15" aria-hidden="true" tabindex="-1"></a><span class="at"> </span><span class="fu">vllm_server_port</span><span class="kw">:</span><span class="at"> </span><span class="dv">8000</span></span>
<span id="cb31-16"><a href="#cb31-16" aria-hidden="true" tabindex="-1"></a><span class="at"> </span><span class="fu">vllm_lora_sync</span><span class="kw">:</span><span class="at"> </span><span class="ch">true</span></span>
<span id="cb31-17"><a href="#cb31-17" aria-hidden="true" tabindex="-1"></a><span class="at"> </span><span class="fu">vllm_sync_interval</span><span class="kw">:</span><span class="at"> </span><span class="dv">5</span></span>
<span id="cb31-18"><a href="#cb31-18" aria-hidden="true" tabindex="-1"></a><span class="co"> # Async GRPO — 3x faster than standard</span></span>
<span id="cb31-19"><a href="#cb31-19" aria-hidden="true" tabindex="-1"></a><span class="at"> </span><span class="fu">use_data_producer</span><span class="kw">:</span><span class="at"> </span><span class="ch">true</span></span>
<span id="cb31-20"><a href="#cb31-20" aria-hidden="true" tabindex="-1"></a><span class="at"> </span><span class="fu">async_prefetch</span><span class="kw">:</span><span class="at"> </span><span class="ch">true</span></span>
<span id="cb31-21"><a href="#cb31-21" aria-hidden="true" tabindex="-1"></a><span class="at"> </span><span class="fu">num_generations</span><span class="kw">:</span><span class="at"> </span><span class="dv">4</span></span>
<span id="cb31-22"><a href="#cb31-22" aria-hidden="true" tabindex="-1"></a><span class="at"> </span><span class="fu">max_completion_length</span><span class="kw">:</span><span class="at"> </span><span class="dv">512</span></span>
<span id="cb31-23"><a href="#cb31-23" aria-hidden="true" tabindex="-1"></a><span class="at"> </span><span class="fu">temperature</span><span class="kw">:</span><span class="at"> </span><span class="fl">0.8</span></span>
<span id="cb31-24"><a href="#cb31-24" aria-hidden="true" tabindex="-1"></a><span class="at"> </span><span class="fu">reward_funcs</span><span class="kw">:</span></span>
<span id="cb31-25"><a href="#cb31-25" aria-hidden="true" tabindex="-1"></a><span class="at"> </span><span class="kw">-</span><span class="at"> axolotl.integrations.nemo_gym.rewards.reward_env</span></span>
<span id="cb31-26"><a href="#cb31-26" aria-hidden="true" tabindex="-1"></a></span>
<span id="cb31-27"><a href="#cb31-27" aria-hidden="true" tabindex="-1"></a><span class="fu">plugins</span><span class="kw">:</span></span>
<span id="cb31-28"><a href="#cb31-28" aria-hidden="true" tabindex="-1"></a><span class="at"> </span><span class="kw">-</span><span class="at"> axolotl.integrations.nemo_gym.NemoGymPlugin</span></span>
<span id="cb31-29"><a href="#cb31-29" aria-hidden="true" tabindex="-1"></a></span>
<span id="cb31-30"><a href="#cb31-30" aria-hidden="true" tabindex="-1"></a><span class="fu">nemo_gym_enabled</span><span class="kw">:</span><span class="at"> </span><span class="ch">true</span></span>
<span id="cb31-31"><a href="#cb31-31" aria-hidden="true" tabindex="-1"></a><span class="fu">nemo_gym_auto_start</span><span class="kw">:</span><span class="at"> </span><span class="ch">false</span></span>
<span id="cb31-32"><a href="#cb31-32" aria-hidden="true" tabindex="-1"></a><span class="fu">nemo_gym_head_port</span><span class="kw">:</span><span class="at"> </span><span class="dv">11000</span></span>
<span id="cb31-33"><a href="#cb31-33" aria-hidden="true" tabindex="-1"></a><span class="fu">nemo_gym_multi_turn</span><span class="kw">:</span><span class="at"> </span><span class="ch">true</span></span>
<span id="cb31-34"><a href="#cb31-34" aria-hidden="true" tabindex="-1"></a><span class="fu">nemo_gym_verify_timeout</span><span class="kw">:</span><span class="at"> </span><span class="dv">120</span></span>
<span id="cb31-35"><a href="#cb31-35" aria-hidden="true" tabindex="-1"></a><span class="fu">nemo_gym_datasets</span><span class="kw">:</span></span>
<span id="cb31-36"><a href="#cb31-36" aria-hidden="true" tabindex="-1"></a><span class="at"> </span><span class="kw">-</span><span class="at"> </span><span class="fu">path</span><span class="kw">:</span><span class="at"> ~/Gym/resources_servers/example_single_tool_call/data/weather_tool_calling.jsonl</span></span>
<span id="cb31-37"><a href="#cb31-37" aria-hidden="true" tabindex="-1"></a><span class="at"> </span><span class="fu">server_name</span><span class="kw">:</span><span class="at"> example_single_tool_call</span></span>
<span id="cb31-38"><a href="#cb31-38" aria-hidden="true" tabindex="-1"></a></span>
<span id="cb31-39"><a href="#cb31-39" aria-hidden="true" tabindex="-1"></a><span class="fu">datasets</span><span class="kw">:</span></span>
<span id="cb31-40"><a href="#cb31-40" aria-hidden="true" tabindex="-1"></a><span class="at"> </span><span class="kw">-</span><span class="at"> </span><span class="fu">path</span><span class="kw">:</span><span class="at"> ~/Gym/resources_servers/example_single_tool_call/data/weather_tool_calling.jsonl</span></span>
<span id="cb31-41"><a href="#cb31-41" aria-hidden="true" tabindex="-1"></a><span class="at"> </span><span class="fu">type</span><span class="kw">:</span><span class="at"> chat_template</span></span>
<span id="cb31-42"><a href="#cb31-42" aria-hidden="true" tabindex="-1"></a><span class="at"> </span><span class="fu">field_messages</span><span class="kw">:</span><span class="at"> responses_create_params.input</span></span>
<span id="cb31-43"><a href="#cb31-43" aria-hidden="true" tabindex="-1"></a><span class="at"> </span><span class="fu">message_field_content</span><span class="kw">:</span><span class="at"> content</span></span>
<span id="cb31-44"><a href="#cb31-44" aria-hidden="true" tabindex="-1"></a><span class="at"> </span><span class="fu">message_field_role</span><span class="kw">:</span><span class="at"> role</span></span>
<span id="cb31-45"><a href="#cb31-45" aria-hidden="true" tabindex="-1"></a></span>
<span id="cb31-46"><a href="#cb31-46" aria-hidden="true" tabindex="-1"></a><span class="fu">vllm</span><span class="kw">:</span></span>
<span id="cb31-47"><a href="#cb31-47" aria-hidden="true" tabindex="-1"></a><span class="at"> </span><span class="fu">gpu_memory_utilization</span><span class="kw">:</span><span class="at"> </span><span class="fl">0.85</span></span>
<span id="cb31-48"><a href="#cb31-48" aria-hidden="true" tabindex="-1"></a><span class="at"> </span><span class="fu">max_model_len</span><span class="kw">:</span><span class="at"> </span><span class="dv">2048</span></span>
<span id="cb31-49"><a href="#cb31-49" aria-hidden="true" tabindex="-1"></a><span class="at"> </span><span class="fu">tensor_parallel_size</span><span class="kw">:</span><span class="at"> </span><span class="dv">1</span></span>
<span id="cb31-50"><a href="#cb31-50" aria-hidden="true" tabindex="-1"></a></span>
<span id="cb31-51"><a href="#cb31-51" aria-hidden="true" tabindex="-1"></a><span class="fu">learning_rate</span><span class="kw">:</span><span class="at"> </span><span class="fl">5e-6</span></span>
<span id="cb31-52"><a href="#cb31-52" aria-hidden="true" tabindex="-1"></a><span class="fu">micro_batch_size</span><span class="kw">:</span><span class="at"> </span><span class="dv">1</span></span>
<span id="cb31-53"><a href="#cb31-53" aria-hidden="true" tabindex="-1"></a><span class="fu">gradient_accumulation_steps</span><span class="kw">:</span><span class="at"> </span><span class="dv">4</span></span>
<span id="cb31-54"><a href="#cb31-54" aria-hidden="true" tabindex="-1"></a><span class="fu">max_steps</span><span class="kw">:</span><span class="at"> </span><span class="dv">30</span></span>
<span id="cb31-55"><a href="#cb31-55" aria-hidden="true" tabindex="-1"></a><span class="fu">gradient_checkpointing</span><span class="kw">:</span><span class="at"> </span><span class="ch">true</span></span>
<span id="cb31-56"><a href="#cb31-56" aria-hidden="true" tabindex="-1"></a><span class="fu">bf16</span><span class="kw">:</span><span class="at"> </span><span class="ch">true</span></span>
<span id="cb31-57"><a href="#cb31-57" aria-hidden="true" tabindex="-1"></a><span class="fu">output_dir</span><span class="kw">:</span><span class="at"> ./outputs/nemo_gym_async</span></span>
<span id="cb31-58"><a href="#cb31-58" aria-hidden="true" tabindex="-1"></a></span>
<span id="cb31-59"><a href="#cb31-59" aria-hidden="true" tabindex="-1"></a><span class="fu">use_wandb</span><span class="kw">:</span><span class="at"> </span><span class="ch">true</span></span>
<span id="cb31-60"><a href="#cb31-60" aria-hidden="true" tabindex="-1"></a><span class="fu">wandb_project</span><span class="kw">:</span><span class="at"> nemo-gym-rl</span></span></code></pre></div><button title="Copy to Clipboard" class="code-copy-button"><i class="bi"></i></button></div>
</section>
<section id="single-turn-training-simplest-no-agent-server-needed" class="level3">
<h3 class="anchored" data-anchor-id="single-turn-training-simplest-no-agent-server-needed">3. Single-Turn Training (Simplest — No Agent Server Needed)</h3>
<p>For environments that only need single-turn verify (math, coding challenges), you don’t need
an agent server. The plugin’s reward function calls <code>/verify</code> directly.</p>
<div class="code-copy-outer-scaffold"><div class="sourceCode" id="cb32"><pre class="sourceCode yaml code-with-copy"><code class="sourceCode yaml"><span id="cb32-1"><a href="#cb32-1" aria-hidden="true" tabindex="-1"></a><span class="fu">base_model</span><span class="kw">:</span><span class="at"> Qwen/Qwen2.5-0.5B-Instruct</span></span>
<span id="cb32-2"><a href="#cb32-2" aria-hidden="true" tabindex="-1"></a><span class="fu">rl</span><span class="kw">:</span><span class="at"> grpo</span></span>
<span id="cb32-3"><a href="#cb32-3" aria-hidden="true" tabindex="-1"></a><span class="fu">chat_template</span><span class="kw">:</span><span class="at"> tokenizer_default</span></span>
<span id="cb32-4"><a href="#cb32-4" aria-hidden="true" tabindex="-1"></a></span>
<span id="cb32-5"><a href="#cb32-5" aria-hidden="true" tabindex="-1"></a><span class="fu">trl</span><span class="kw">:</span></span>
<span id="cb32-6"><a href="#cb32-6" aria-hidden="true" tabindex="-1"></a><span class="at"> </span><span class="fu">use_vllm</span><span class="kw">:</span><span class="at"> </span><span class="ch">true</span></span>
<span id="cb32-7"><a href="#cb32-7" aria-hidden="true" tabindex="-1"></a><span class="at"> </span><span class="fu">vllm_mode</span><span class="kw">:</span><span class="at"> colocate</span></span>
<span id="cb32-8"><a href="#cb32-8" aria-hidden="true" tabindex="-1"></a><span class="at"> </span><span class="fu">vllm_enable_sleep_mode</span><span class="kw">:</span><span class="at"> </span><span class="ch">false</span></span>
<span id="cb32-9"><a href="#cb32-9" aria-hidden="true" tabindex="-1"></a><span class="at"> </span><span class="fu">num_generations</span><span class="kw">:</span><span class="at"> </span><span class="dv">8</span></span>
<span id="cb32-10"><a href="#cb32-10" aria-hidden="true" tabindex="-1"></a><span class="at"> </span><span class="fu">max_completion_length</span><span class="kw">:</span><span class="at"> </span><span class="dv">128</span></span>
<span id="cb32-11"><a href="#cb32-11" aria-hidden="true" tabindex="-1"></a><span class="at"> </span><span class="fu">temperature</span><span class="kw">:</span><span class="at"> </span><span class="fl">0.9</span></span>
<span id="cb32-12"><a href="#cb32-12" aria-hidden="true" tabindex="-1"></a><span class="at"> </span><span class="fu">reward_funcs</span><span class="kw">:</span></span>
<span id="cb32-13"><a href="#cb32-13" aria-hidden="true" tabindex="-1"></a><span class="at"> </span><span class="kw">-</span><span class="at"> axolotl.integrations.nemo_gym.rewards.reward_nemo_gym_verify</span></span>
<span id="cb32-14"><a href="#cb32-14" aria-hidden="true" tabindex="-1"></a></span>
<span id="cb32-15"><a href="#cb32-15" aria-hidden="true" tabindex="-1"></a><span class="fu">plugins</span><span class="kw">:</span></span>
<span id="cb32-16"><a href="#cb32-16" aria-hidden="true" tabindex="-1"></a><span class="at"> </span><span class="kw">-</span><span class="at"> axolotl.integrations.nemo_gym.NemoGymPlugin</span></span>
<span id="cb32-17"><a href="#cb32-17" aria-hidden="true" tabindex="-1"></a></span>
<span id="cb32-18"><a href="#cb32-18" aria-hidden="true" tabindex="-1"></a><span class="fu">nemo_gym_enabled</span><span class="kw">:</span><span class="at"> </span><span class="ch">true</span></span>
<span id="cb32-19"><a href="#cb32-19" aria-hidden="true" tabindex="-1"></a><span class="fu">nemo_gym_auto_start</span><span class="kw">:</span><span class="at"> </span><span class="ch">false</span></span>
<span id="cb32-20"><a href="#cb32-20" aria-hidden="true" tabindex="-1"></a><span class="fu">nemo_gym_head_port</span><span class="kw">:</span><span class="at"> </span><span class="dv">11000</span></span>
<span id="cb32-21"><a href="#cb32-21" aria-hidden="true" tabindex="-1"></a><span class="fu">nemo_gym_datasets</span><span class="kw">:</span></span>
<span id="cb32-22"><a href="#cb32-22" aria-hidden="true" tabindex="-1"></a><span class="at"> </span><span class="kw">-</span><span class="at"> </span><span class="fu">path</span><span class="kw">:</span><span class="at"> ~/Gym/resources_servers/reasoning_gym/data/train_basic_arithmetic.jsonl</span></span>
<span id="cb32-23"><a href="#cb32-23" aria-hidden="true" tabindex="-1"></a><span class="at"> </span><span class="fu">server_name</span><span class="kw">:</span><span class="at"> reasoning_gym</span></span>
<span id="cb32-24"><a href="#cb32-24" aria-hidden="true" tabindex="-1"></a></span>
<span id="cb32-25"><a href="#cb32-25" aria-hidden="true" tabindex="-1"></a><span class="fu">datasets</span><span class="kw">:</span></span>
<span id="cb32-26"><a href="#cb32-26" aria-hidden="true" tabindex="-1"></a><span class="at"> </span><span class="kw">-</span><span class="at"> </span><span class="fu">path</span><span class="kw">:</span><span class="at"> ~/Gym/resources_servers/reasoning_gym/data/train_basic_arithmetic.jsonl</span></span>
<span id="cb32-27"><a href="#cb32-27" aria-hidden="true" tabindex="-1"></a><span class="at"> </span><span class="fu">type</span><span class="kw">:</span><span class="at"> chat_template</span></span>
<span id="cb32-28"><a href="#cb32-28" aria-hidden="true" tabindex="-1"></a><span class="at"> </span><span class="fu">field_messages</span><span class="kw">:</span><span class="at"> responses_create_params.input</span></span>
<span id="cb32-29"><a href="#cb32-29" aria-hidden="true" tabindex="-1"></a><span class="at"> </span><span class="fu">message_field_content</span><span class="kw">:</span><span class="at"> content</span></span>
<span id="cb32-30"><a href="#cb32-30" aria-hidden="true" tabindex="-1"></a><span class="at"> </span><span class="fu">message_field_role</span><span class="kw">:</span><span class="at"> role</span></span>
<span id="cb32-31"><a href="#cb32-31" aria-hidden="true" tabindex="-1"></a></span>
<span id="cb32-32"><a href="#cb32-32" aria-hidden="true" tabindex="-1"></a><span class="fu">vllm</span><span class="kw">:</span></span>
<span id="cb32-33"><a href="#cb32-33" aria-hidden="true" tabindex="-1"></a><span class="at"> </span><span class="fu">gpu_memory_utilization</span><span class="kw">:</span><span class="at"> </span><span class="fl">0.3</span></span>
<span id="cb32-34"><a href="#cb32-34" aria-hidden="true" tabindex="-1"></a><span class="at"> </span><span class="fu">max_model_len</span><span class="kw">:</span><span class="at"> </span><span class="dv">512</span></span>
<span id="cb32-35"><a href="#cb32-35" aria-hidden="true" tabindex="-1"></a><span class="at"> </span><span class="fu">tensor_parallel_size</span><span class="kw">:</span><span class="at"> </span><span class="dv">1</span></span>
<span id="cb32-36"><a href="#cb32-36" aria-hidden="true" tabindex="-1"></a></span>
<span id="cb32-37"><a href="#cb32-37" aria-hidden="true" tabindex="-1"></a><span class="fu">learning_rate</span><span class="kw">:</span><span class="at"> </span><span class="fl">1e-5</span></span>
<span id="cb32-38"><a href="#cb32-38" aria-hidden="true" tabindex="-1"></a><span class="fu">micro_batch_size</span><span class="kw">:</span><span class="at"> </span><span class="dv">4</span></span>
<span id="cb32-39"><a href="#cb32-39" aria-hidden="true" tabindex="-1"></a><span class="fu">gradient_accumulation_steps</span><span class="kw">:</span><span class="at"> </span><span class="dv">2</span></span>
<span id="cb32-40"><a href="#cb32-40" aria-hidden="true" tabindex="-1"></a><span class="fu">max_steps</span><span class="kw">:</span><span class="at"> </span><span class="dv">50</span></span>
<span id="cb32-41"><a href="#cb32-41" aria-hidden="true" tabindex="-1"></a><span class="fu">output_dir</span><span class="kw">:</span><span class="at"> ./outputs/nemo_gym_arithmetic</span></span></code></pre></div><button title="Copy to Clipboard" class="code-copy-button"><i class="bi"></i></button></div>
<p>Only needs <code>ng_run</code> with resource servers (no agent config):</p>
<div class="code-copy-outer-scaffold"><div class="sourceCode" id="cb33"><pre class="sourceCode bash code-with-copy"><code class="sourceCode bash"><span id="cb33-1"><a href="#cb33-1" aria-hidden="true" tabindex="-1"></a><span class="bu">cd</span> ~/Gym <span class="kw">&amp;&amp;</span> <span class="ex">ng_run</span> <span class="st">"+config_paths=[resources_servers/reasoning_gym/configs/resources_only.yaml]"</span> <span class="st">"+skip_venv_if_present=true"</span></span></code></pre></div><button title="Copy to Clipboard" class="code-copy-button"><i class="bi"></i></button></div>
</section>
<section id="how-it-works-2" class="level3">
<h3 class="anchored" data-anchor-id="how-it-works-2">How It Works</h3>
</section>
<section id="single-turn" class="level3">
<h3 class="anchored" data-anchor-id="single-turn">Single-Turn</h3>
<pre class="text"><code>axolotl train → GRPO Trainer generates completions
→ NeMo Gym plugin reward_fn calls POST /verify on resource server
→ reward flows back to GRPO for advantage computation</code></pre>
</section>
<section id="multi-turn-agent-run" class="level3">
<h3 class="anchored" data-anchor-id="multi-turn-agent-run">Multi-Turn (Agent /run)</h3>
<pre class="text"><code>┌─────────────┐ ┌──────────────┐ ┌──────────────────┐
│ axolotl │ │ NeMo Gym │────▶│ vLLM OpenAI │
│ train │────▶│ Agent /run │◀────│ Server (GPU 0) │
│ (GPU 1) │ │ │ │ /v1/completions │
└─────────────┘ └──────┬───────┘ └──────────────────┘
│
▼
┌──────────────┐
│ Resource │
│ Server │
│ (tools + │
│ verify) │
└─────────────┘</code></pre>
<p>The agent server orchestrates the entire multi-turn loop:
1. Calls our vLLM server for model generation
2. Parses tool calls from model output
3. Executes tools against resource servers
4. Feeds tool results back to the model
5. Repeats until done, then calls /verify for reward
6. Returns token IDs + logprobs + reward to our rollout_func</p>
</section>
<section id="data-producer-architecture-multi-turn" class="level3">
<h3 class="anchored" data-anchor-id="data-producer-architecture-multi-turn">Data Producer Architecture (Multi-Turn)</h3>
<p>When <code>nemo_gym_multi_turn: true</code>, the plugin automatically forces <code>use_data_producer: true</code>
which selects the <code>AxolotlAsyncGRPOTrainer</code>. The plugin then swaps the trainer’s data
producer with <code>NemoGymDataProducer</code>, which:</p>
<ol type="1">
<li>Gets a prompt batch from the dataset iterator</li>
<li>Expands by <code>num_generations</code> (one agent call per rollout)</li>
<li>Calls NeMo Gym agents via async HTTP (<code>aiohttp.gather</code>)</li>
<li>Parses responses into padded tensors (<code>RolloutDataset</code>)</li>
<li>Returns with <code>_pending_policy_logps=True</code> for deferred scoring</li>
</ol>
<p>The main thread then runs <code>_compute_deferred_scores()</code> which:
- Computes <strong>policy logprobs</strong> on the training model (GPU forward pass)
- Computes <strong>IS correction</strong> using agent’s sampling logprobs vs training model logprobs
- Computes advantages with group-level normalization
- All downstream features work: replay buffer, re-roll, streaming, zero-adv skip</p>
<p>With <code>async_prefetch: true</code>, the data producer runs in a background thread — giving ~3x
speedup as generation and training overlap. With <code>async_prefetch: false</code>, it runs
synchronously on the main thread (still uses the data producer protocol).</p>
</section>
<section id="weight-sync-lora-mode" class="level3">
<h3 class="anchored" data-anchor-id="weight-sync-lora-mode">Weight Sync (LoRA Mode)</h3>
<p>With <code>vllm_lora_sync: true</code>, the plugin (or async trainer) replaces NCCL-based weight
sync with filesystem + HTTP:</p>
<ol type="1">
<li><code>accelerator.get_state_dict()</code> gathers LoRA weights from all ranks</li>
<li>Rank 0 saves adapter to <code>/tmp/lora_sync_*/vN/</code></li>
<li>Rank 0 POSTs to <code>/set_lora_adapter/</code> on vLLM server</li>
<li>vLLM loads adapter natively via Punica kernels</li>
<li>Only ~40MB transferred (vs multiple GBs for full model weights)</li>
</ol>
</section>
<section id="multi-environment-support" class="level3">
<h3 class="anchored" data-anchor-id="multi-environment-support">Multi-Environment Support</h3>
<p>Datasets support per-row environment routing via <code>agent_ref</code>:</p>
<div class="code-copy-outer-scaffold"><div class="sourceCode" id="cb36"><pre class="sourceCode jsonl code-with-copy"><code class="sourceCode json"><span id="cb36-1"><a href="#cb36-1" aria-hidden="true" tabindex="-1"></a><span class="fu">{</span><span class="dt">"agent_ref"</span><span class="fu">:</span> <span class="fu">{</span><span class="dt">"name"</span><span class="fu">:</span> <span class="st">"reasoning_gym"</span><span class="fu">},</span> <span class="dt">"responses_create_params"</span><span class="fu">:</span> <span class="fu">{</span><span class="er">...</span><span class="fu">}}</span></span>
<span id="cb36-2"><a href="#cb36-2" aria-hidden="true" tabindex="-1"></a><span class="fu">{</span><span class="dt">"agent_ref"</span><span class="fu">:</span> <span class="fu">{</span><span class="dt">"name"</span><span class="fu">:</span> <span class="st">"instruction_following"</span><span class="fu">},</span> <span class="dt">"responses_create_params"</span><span class="fu">:</span> <span class="fu">{</span><span class="er">...</span><span class="fu">}}</span></span></code></pre></div><button title="Copy to Clipboard" class="code-copy-button"><i class="bi"></i></button></div>
<p>Or use the simpler per-dataset routing:</p>
<div class="code-copy-outer-scaffold"><div class="sourceCode" id="cb37"><pre class="sourceCode yaml code-with-copy"><code class="sourceCode yaml"><span id="cb37-1"><a href="#cb37-1" aria-hidden="true" tabindex="-1"></a><span class="fu">nemo_gym_datasets</span><span class="kw">:</span></span>
<span id="cb37-2"><a href="#cb37-2" aria-hidden="true" tabindex="-1"></a><span class="at"> </span><span class="kw">-</span><span class="at"> </span><span class="fu">path</span><span class="kw">:</span><span class="at"> reasoning_data.jsonl</span></span>
<span id="cb37-3"><a href="#cb37-3" aria-hidden="true" tabindex="-1"></a><span class="at"> </span><span class="fu">server_name</span><span class="kw">:</span><span class="at"> reasoning_gym</span></span>
<span id="cb37-4"><a href="#cb37-4" aria-hidden="true" tabindex="-1"></a><span class="at"> </span><span class="kw">-</span><span class="at"> </span><span class="fu">path</span><span class="kw">:</span><span class="at"> tool_data.jsonl</span></span>
<span id="cb37-5"><a href="#cb37-5" aria-hidden="true" tabindex="-1"></a><span class="at"> </span><span class="fu">server_name</span><span class="kw">:</span><span class="at"> example_single_tool_call</span></span></code></pre></div><button title="Copy to Clipboard" class="code-copy-button"><i class="bi"></i></button></div>
</section>
<section id="configuration-reference" class="level3">
<h3 class="anchored" data-anchor-id="configuration-reference">Configuration Reference</h3>
<table class="caption-top table">
<colgroup>
<col style="width: 28%">
<col style="width: 15%">
<col style="width: 23%">
<col style="width: 33%">
</colgroup>
<thead>
<tr class="header">
<th>Parameter</th>
<th>Type</th>
<th>Default</th>
<th>Description</th>
</tr>
</thead>
<tbody>
<tr class="odd">
<td><code>nemo_gym_enabled</code></td>
<td>bool</td>
<td><code>null</code></td>
<td>Enable the NeMo Gym integration</td>
</tr>
<tr class="even">
<td><code>nemo_gym_dir</code></td>
<td>str</td>
<td><code>~/Gym</code></td>
<td>Path to NeMo Gym repo</td>
</tr>
<tr class="odd">
<td><code>nemo_gym_auto_clone</code></td>
<td>bool</td>
<td><code>true</code></td>
<td>Auto-clone NeMo Gym repo if missing</td>
</tr>
<tr class="even">
<td><code>nemo_gym_auto_start</code></td>
<td>bool</td>
<td><code>true</code></td>
<td>Auto-start resource servers</td>
</tr>
<tr class="odd">
<td><code>nemo_gym_config_paths</code></td>
<td>list[str]</td>
<td>—</td>
<td>Server config YAMLs (relative to gym_dir)</td>
</tr>
<tr class="even">
<td><code>nemo_gym_datasets</code></td>
<td>list[dict]</td>
<td>required</td>
<td>Dataset configs with <code>path</code> and optional <code>server_name</code></td>
</tr>
<tr class="odd">
<td><code>nemo_gym_head_port</code></td>
<td>int</td>
<td><code>11000</code></td>
<td>Head server port</td>
</tr>
<tr class="even">
<td><code>nemo_gym_server_timeout</code></td>
<td>int</td>
<td><code>360</code></td>
<td>Server startup timeout (seconds)</td>
</tr>
<tr class="odd">
<td><code>nemo_gym_verify_timeout</code></td>
<td>int</td>
<td><code>30</code></td>
<td>Per-request timeout (seconds)</td>
</tr>
<tr class="even">
<td><code>nemo_gym_multi_turn</code></td>
<td>bool</td>
<td><code>false</code></td>
<td>Enable multi-turn via agent /run</td>
</tr>
</tbody>
</table>
</section>
<section id="dataset-jsonl-format" class="level3">
<h3 class="anchored" data-anchor-id="dataset-jsonl-format">Dataset JSONL Format</h3>
<p>Each line must have <code>responses_create_params</code> with <code>input</code> messages:</p>
<div class="code-copy-outer-scaffold"><div class="sourceCode" id="cb38"><pre class="sourceCode json code-with-copy"><code class="sourceCode json"><span id="cb38-1"><a href="#cb38-1" aria-hidden="true" tabindex="-1"></a><span class="fu">{</span></span>
<span id="cb38-2"><a href="#cb38-2" aria-hidden="true" tabindex="-1"></a> <span class="dt">"responses_create_params"</span><span class="fu">:</span> <span class="fu">{</span></span>
<span id="cb38-3"><a href="#cb38-3" aria-hidden="true" tabindex="-1"></a> <span class="dt">"input"</span><span class="fu">:</span> <span class="ot">[</span><span class="fu">{</span><span class="dt">"role"</span><span class="fu">:</span> <span class="st">"user"</span><span class="fu">,</span> <span class="dt">"content"</span><span class="fu">:</span> <span class="st">"What's the weather in SF?"</span><span class="fu">}</span><span class="ot">]</span><span class="fu">,</span></span>
<span id="cb38-4"><a href="#cb38-4" aria-hidden="true" tabindex="-1"></a> <span class="dt">"tools"</span><span class="fu">:</span> <span class="ot">[</span><span class="fu">{</span><span class="dt">"name"</span><span class="fu">:</span> <span class="st">"get_weather"</span><span class="fu">,</span> <span class="dt">"type"</span><span class="fu">:</span> <span class="st">"function"</span><span class="fu">,</span> <span class="dt">"strict"</span><span class="fu">:</span> <span class="kw">true</span><span class="fu">,</span> <span class="dt">"parameters"</span><span class="fu">:</span> <span class="fu">{</span><span class="er">...</span><span class="fu">}}</span><span class="ot">]</span></span>
<span id="cb38-5"><a href="#cb38-5" aria-hidden="true" tabindex="-1"></a> <span class="fu">}</span></span>
<span id="cb38-6"><a href="#cb38-6" aria-hidden="true" tabindex="-1"></a><span class="fu">}</span></span></code></pre></div><button title="Copy to Clipboard" class="code-copy-button"><i class="bi"></i></button></div>
<p>For multi-turn agent routing, include <code>agent_ref</code>:</p>
<div class="code-copy-outer-scaffold"><div class="sourceCode" id="cb39"><pre class="sourceCode json code-with-copy"><code class="sourceCode json"><span id="cb39-1"><a href="#cb39-1" aria-hidden="true" tabindex="-1"></a><span class="fu">{</span><span class="dt">"agent_ref"</span><span class="fu">:</span> <span class="fu">{</span><span class="dt">"name"</span><span class="fu">:</span> <span class="st">"my_agent"</span><span class="fu">},</span> <span class="dt">"responses_create_params"</span><span class="fu">:</span> <span class="fu">{</span><span class="er">...</span><span class="fu">}}</span></span></code></pre></div><button title="Copy to Clipboard" class="code-copy-button"><i class="bi"></i></button></div>
<p>Note: Tool definitions MUST include <code>"strict": true</code> and <code>"additionalProperties": false</code> for NeMo Gym agent compatibility.</p>
</section>
<section id="reward-functions" class="level3">
<h3 class="anchored" data-anchor-id="reward-functions">Reward Functions</h3>
<p>The plugin provides two built-in reward functions — no user code needed:</p>
<div class="code-copy-outer-scaffold"><div class="sourceCode" id="cb40"><pre class="sourceCode yaml code-with-copy"><code class="sourceCode yaml"><span id="cb40-1"><a href="#cb40-1" aria-hidden="true" tabindex="-1"></a><span class="fu">trl</span><span class="kw">:</span></span>
<span id="cb40-2"><a href="#cb40-2" aria-hidden="true" tabindex="-1"></a><span class="at"> </span><span class="fu">reward_funcs</span><span class="kw">:</span></span>
<span id="cb40-3"><a href="#cb40-3" aria-hidden="true" tabindex="-1"></a><span class="co"> # Multi-turn (nemo_gym_multi_turn: true):</span></span>
<span id="cb40-4"><a href="#cb40-4" aria-hidden="true" tabindex="-1"></a><span class="co"> # Passthrough — agent /run already computed the reward</span></span>
<span id="cb40-5"><a href="#cb40-5" aria-hidden="true" tabindex="-1"></a><span class="at"> </span><span class="kw">-</span><span class="at"> axolotl.integrations.nemo_gym.rewards.reward_env</span></span>
<span id="cb40-6"><a href="#cb40-6" aria-hidden="true" tabindex="-1"></a></span>
<span id="cb40-7"><a href="#cb40-7" aria-hidden="true" tabindex="-1"></a><span class="co"> # Single-turn (nemo_gym_multi_turn: false):</span></span>
<span id="cb40-8"><a href="#cb40-8" aria-hidden="true" tabindex="-1"></a><span class="co"> # Calls /verify endpoints on NeMo Gym resource servers</span></span>
<span id="cb40-9"><a href="#cb40-9" aria-hidden="true" tabindex="-1"></a><span class="at"> </span><span class="kw">-</span><span class="at"> axolotl.integrations.nemo_gym.rewards.reward_nemo_gym_verify</span></span></code></pre></div><button title="Copy to Clipboard" class="code-copy-button"><i class="bi"></i></button></div>
<p>Both are also importable from Python:</p>
<div class="code-copy-outer-scaffold"><div class="sourceCode" id="cb41"><pre class="sourceCode python code-with-copy"><code class="sourceCode python"><span id="cb41-1"><a href="#cb41-1" aria-hidden="true" tabindex="-1"></a><span class="im">from</span> axolotl.integrations.nemo_gym <span class="im">import</span> reward_env, reward_nemo_gym_verify</span></code></pre></div><button title="Copy to Clipboard" class="code-copy-button"><i class="bi"></i></button></div>
</section>
<section id="known-issues-troubleshooting" class="level3">
<h3 class="anchored" data-anchor-id="known-issues-troubleshooting">Known Issues / Troubleshooting</h3>
</section>
<section id="nemo-gym-server-setup" class="level3">
<h3 class="anchored" data-anchor-id="nemo-gym-server-setup">NeMo Gym Server Setup</h3>
<ul>
<li><strong>pycosat build failure</strong>: <code>CFLAGS="" uv pip install pycosat --no-build-isolation</code></li>
<li><strong>Ray version mismatch</strong>: Pin <code>ray[default]==2.52.1</code> in all server venvs</li>
<li><strong>Pre-build venvs</strong>: <code>ng_run</code> creates per-server venvs via Ray. Pre-build them and use <code>+skip_venv_if_present=true</code></li>
<li><strong>Tool <code>strict</code> field required</strong>: Agent server validates tool definitions require <code>strict: true</code></li>
</ul>
</section>
<section id="vllm-weight-sync" class="level3">
<h3 class="anchored" data-anchor-id="vllm-weight-sync">vLLM / Weight Sync</h3>
<ul>
<li><p><strong>Start vLLM with LoRA + tool calling + runtime loading</strong>:</p>
<div class="code-copy-outer-scaffold"><div class="sourceCode" id="cb42"><pre class="sourceCode bash code-with-copy"><code class="sourceCode bash"><span id="cb42-1"><a href="#cb42-1" aria-hidden="true" tabindex="-1"></a><span class="va">VLLM_ALLOW_RUNTIME_LORA_UPDATING</span><span class="op">=</span>1 <span class="dt">\</span></span>
<span id="cb42-2"><a href="#cb42-2" aria-hidden="true" tabindex="-1"></a><span class="va">CUDA_VISIBLE_DEVICES</span><span class="op">=</span>0 <span class="ex">python</span> <span class="at">-m</span> vllm.entrypoints.openai.api_server <span class="dt">\</span></span>
<span id="cb42-3"><a href="#cb42-3" aria-hidden="true" tabindex="-1"></a> <span class="at">--model</span> Qwen/Qwen3-4B-Instruct-2507 <span class="dt">\</span></span>
<span id="cb42-4"><a href="#cb42-4" aria-hidden="true" tabindex="-1"></a> <span class="at">--max-model-len</span> 4096 <span class="dt">\</span></span>
<span id="cb42-5"><a href="#cb42-5" aria-hidden="true" tabindex="-1"></a> <span class="at">--gpu-memory-utilization</span> 0.7 <span class="dt">\</span></span>
<span id="cb42-6"><a href="#cb42-6" aria-hidden="true" tabindex="-1"></a> <span class="at">--enable-lora</span> <span class="at">--max-lora-rank</span> 64 <span class="dt">\</span></span>
<span id="cb42-7"><a href="#cb42-7" aria-hidden="true" tabindex="-1"></a> <span class="at">--enable-auto-tool-choice</span> <span class="at">--tool-call-parser</span> hermes</span></code></pre></div><button title="Copy to Clipboard" class="code-copy-button"><i class="bi"></i></button></div></li>
<li><p><strong><code>VLLM_ALLOW_RUNTIME_LORA_UPDATING=1</code></strong>: Required for <code>vllm_lora_sync: true</code>. Without it, vLLM won’t expose the <code>/v1/load_lora_adapter</code> endpoint and weight sync will fail silently. The plugin warns if this endpoint is missing.</p></li>
<li><p><strong><code>--enable-lora</code></strong>: Enables LoRA adapter support in vLLM</p></li>
<li><p><strong><code>--enable-auto-tool-choice --tool-call-parser hermes</code></strong>: Required for Qwen3 tool calling</p></li>
<li><p><strong><code>max_model_len</code> must be &gt; <code>max_completion_length</code></strong>: Leave room for prompt tokens (~200). If equal, the NeMo Gym model proxy gets a 400 error and returns empty completions.</p></li>
<li><p><strong><code>CUDA_HOME</code> required</strong>: DeepSpeed import needs it for the nvcc shim</p></li>
<li><p><strong>NCCL weight sync broken with vLLM 0.17</strong>: Use <code>vllm_lora_sync: true</code> (filesystem + HTTP via <code>/v1/load_lora_adapter</code>)</p></li>
</ul>
</section>
<section id="multi-turn" class="level3">
<h3 class="anchored" data-anchor-id="multi-turn">Multi-Turn</h3>
<ul>
<li><strong>Agent server required</strong>: Multi-turn delegates to NeMo Gym’s agent server <code>/run</code> endpoint. Without an agent, the plugin falls back to single-turn <code>/verify</code></li>
<li><strong>Model server proxy</strong>: NeMo Gym needs a <code>responses_api_models</code> server that proxies to your vLLM. See the agent config example above</li>
</ul>
</section>
<section id="fsdp2" class="level3">
<h3 class="anchored" data-anchor-id="fsdp2">FSDP2</h3>
<ul>
<li>Validated on 2 GPUs with single-turn + LoRA</li>
<li>Async field filtering: The builder automatically filters async-only config fields when using the standard GRPO trainer</li>
</ul>
</section>
<section id="comparison-with-other-integrations" class="level3">
<h3 class="anchored" data-anchor-id="comparison-with-other-integrations">Comparison with Other Integrations</h3>
<table class="caption-top table">
<colgroup>
<col style="width: 13%">
<col style="width: 28%">
<col style="width: 28%">
<col style="width: 28%">
</colgroup>
<thead>
<tr class="header">
<th>Feature</th>
<th>Axolotl + NeMo Gym</th>
<th>Unsloth + NeMo Gym</th>
<th>NeMo RL (native)</th>
</tr>
</thead>
<tbody>
<tr class="odd">
<td>Server management</td>
<td>Automatic</td>
<td>Manual (notebook)</td>
<td>Built-in</td>
</tr>
<tr class="even">
<td>Multi-environment</td>
<td>Per-row routing</td>
<td>Manual code</td>
<td>YAML config</td>
</tr>
<tr class="odd">
<td>Multi-turn / tool use</td>
<td>Agent /run delegation</td>
<td>No</td>
<td>Agent /run (Ray)</td>
</tr>
<tr class="even">
<td>Async GRPO (3x speedup)</td>
<td>Yes</td>
<td>No</td>
<td>Yes</td>
</tr>
<tr class="odd">
<td>LoRA sync</td>
<td>Filesystem + HTTP</td>
<td>N/A</td>
<td>NCCL</td>
</tr>
<tr class="even">
<td>Multi-GPU (FSDP2)</td>
<td>Yes</td>
<td>No</td>
<td>Yes (Ray)</td>
</tr>
<tr class="odd">
<td>Config-driven</td>
<td>Yes</td>
<td>No (code)</td>
<td>Yes</td>
</tr>
</tbody>
</table>
<p>Please see reference <a href="https://github.com/axolotl-ai-cloud/axolotl/tree/main/src/axolotl/integrations/nemo_gym">here</a></p>
</section>
</section>
<section id="spectrum" class="level2">
<h2 class="anchored" data-anchor-id="spectrum">Spectrum</h2>
<p>by Eric Hartford, Lucas Atkins, Fernando Fernandes, David Golchinfar</p>
<p>This plugin contains code to freeze the bottom fraction of modules in a model, based on the Signal-to-Noise Ratio (SNR).</p>
<p>See https://github.com/cognitivecomputations/spectrum</p>
<section id="overview-1" class="level3">
<h3 class="anchored" data-anchor-id="overview-1">Overview</h3>
<p>Spectrum is a tool for scanning and evaluating the Signal-to-Noise Ratio (SNR) of layers in large language models.
By identifying the top n% of layers with the highest SNR, you can optimize training efficiency.</p>
</section>
<section id="usage-7" class="level3">
<h3 class="anchored" data-anchor-id="usage-7">Usage</h3>
<div class="code-copy-outer-scaffold"><div class="sourceCode" id="cb43"><pre class="sourceCode yaml code-with-copy"><code class="sourceCode yaml"><span id="cb43-1"><a href="#cb43-1" aria-hidden="true" tabindex="-1"></a><span class="fu">plugins</span><span class="kw">:</span></span>
<span id="cb43-2"><a href="#cb43-2" aria-hidden="true" tabindex="-1"></a><span class="at"> </span><span class="kw">-</span><span class="at"> axolotl.integrations.spectrum.SpectrumPlugin</span></span>
<span id="cb43-3"><a href="#cb43-3" aria-hidden="true" tabindex="-1"></a></span>
<span id="cb43-4"><a href="#cb43-4" aria-hidden="true" tabindex="-1"></a><span class="fu">spectrum_top_fraction</span><span class="kw">:</span><span class="at"> </span><span class="fl">0.5</span></span>
<span id="cb43-5"><a href="#cb43-5" aria-hidden="true" tabindex="-1"></a><span class="fu">spectrum_model_name</span><span class="kw">:</span><span class="at"> meta-llama/Meta-Llama-3.1-8B</span></span></code></pre></div><button title="Copy to Clipboard" class="code-copy-button"><i class="bi"></i></button></div>
</section>
<section id="citation-4" class="level3">
<h3 class="anchored" data-anchor-id="citation-4">Citation</h3>
<div class="code-copy-outer-scaffold"><div class="sourceCode" id="cb44"><pre class="sourceCode bib code-with-copy"><code class="sourceCode bibtex"><span id="cb44-1"><a href="#cb44-1" aria-hidden="true" tabindex="-1"></a><span class="va">@misc</span>{<span class="ot">hartford2024spectrumtargetedtrainingsignal</span>,</span>
<span id="cb44-2"><a href="#cb44-2" aria-hidden="true" tabindex="-1"></a> <span class="dt">title</span>={Spectrum: Targeted Training on Signal to Noise Ratio},</span>
<span id="cb44-3"><a href="#cb44-3" aria-hidden="true" tabindex="-1"></a> <span class="dt">author</span>={Eric Hartford and Lucas Atkins and Fernando Fernandes Neto and David Golchinfar},</span>
<span id="cb44-4"><a href="#cb44-4" aria-hidden="true" tabindex="-1"></a> <span class="dt">year</span>={2024},</span>
<span id="cb44-5"><a href="#cb44-5" aria-hidden="true" tabindex="-1"></a> <span class="dt">eprint</span>={2406.06623},</span>
<span id="cb44-6"><a href="#cb44-6" aria-hidden="true" tabindex="-1"></a> <span class="dt">archivePrefix</span>={arXiv},</span>
<span id="cb44-7"><a href="#cb44-7" aria-hidden="true" tabindex="-1"></a> <span class="dt">primaryClass</span>={cs.LG},</span>
<span id="cb44-8"><a href="#cb44-8" aria-hidden="true" tabindex="-1"></a> <span class="dt">url</span>={https://arxiv.org/abs/2406.06623},</span>
<span id="cb44-9"><a href="#cb44-9" aria-hidden="true" tabindex="-1"></a>}</span></code></pre></div><button title="Copy to Clipboard" class="code-copy-button"><i class="bi"></i></button></div>
<p>Please see reference <a href="https://github.com/axolotl-ai-cloud/axolotl/tree/main/src/axolotl/integrations/spectrum">here</a></p>
</section>
</section>
<section id="swanlab-integration-for-axolotl" class="level2">
<h2 class="anchored" data-anchor-id="swanlab-integration-for-axolotl">SwanLab Integration for Axolotl</h2>
<p>SwanLab is an open-source, lightweight AI experiment tracking and visualization tool that provides a platform for tracking, recording, comparing, and collaborating on experiments.</p>
<p>This integration enables seamless experiment tracking and visualization of Axolotl training runs using SwanLab.</p>
<section id="features" class="level3">
<h3 class="anchored" data-anchor-id="features">Features</h3>
<ul>
<li>📊 <strong>Automatic Metrics Logging</strong>: Training loss, learning rate, and other metrics are automatically logged</li>
<li>🎯 <strong>Hyperparameter Tracking</strong>: Model configuration and training parameters are tracked</li>
<li>📈 <strong>Real-time Visualization</strong>: Monitor training progress in real-time through SwanLab dashboard</li>
<li>☁️ <strong>Cloud &amp; Local Support</strong>: Works in both cloud-synced and offline modes</li>
<li>🔄 <strong>Experiment Comparison</strong>: Compare multiple training runs easily</li>
<li>🤝 <strong>Team Collaboration</strong>: Share experiments with team members</li>
<li>🎭 <strong>RLHF Completion Logging</strong>: Automatically log model outputs during DPO/KTO/ORPO/GRPO training for qualitative analysis</li>
<li>⚡ <strong>Performance Profiling</strong>: Built-in profiling decorators to measure and optimize training performance</li>
<li>🔔 <strong>Lark Notifications</strong>: Send real-time training updates to team chat (Feishu/Lark integration)</li>
</ul>
</section>
<section id="installation-2" class="level3">
<h3 class="anchored" data-anchor-id="installation-2">Installation</h3>
<div class="code-copy-outer-scaffold"><div class="sourceCode" id="cb45"><pre class="sourceCode bash code-with-copy"><code class="sourceCode bash"><span id="cb45-1"><a href="#cb45-1" aria-hidden="true" tabindex="-1"></a><span class="ex">pip</span> install swanlab</span></code></pre></div><button title="Copy to Clipboard" class="code-copy-button"><i class="bi"></i></button></div>
</section>
<section id="quick-start-1" class="level3">
<h3 class="anchored" data-anchor-id="quick-start-1">Quick Start</h3>
</section>
<section id="register-for-swanlab-optional-for-cloud-mode" class="level3">
<h3 class="anchored" data-anchor-id="register-for-swanlab-optional-for-cloud-mode">1. Register for SwanLab (Optional for cloud mode)</h3>
<p>If you want to use cloud sync features, register at <a href="https://swanlab.cn">https://swanlab.cn</a> to get your API key.</p>
</section>
<section id="configure-axolotl-config-file" class="level3">
<h3 class="anchored" data-anchor-id="configure-axolotl-config-file">2. Configure Axolotl Config File</h3>
<p>Add SwanLab configuration to your Axolotl YAML config:</p>
<div class="code-copy-outer-scaffold"><div class="sourceCode" id="cb46"><pre class="sourceCode yaml code-with-copy"><code class="sourceCode yaml"><span id="cb46-1"><a href="#cb46-1" aria-hidden="true" tabindex="-1"></a><span class="fu">plugins</span><span class="kw">:</span></span>
<span id="cb46-2"><a href="#cb46-2" aria-hidden="true" tabindex="-1"></a><span class="at"> </span><span class="kw">-</span><span class="at"> axolotl.integrations.swanlab.SwanLabPlugin</span></span>
<span id="cb46-3"><a href="#cb46-3" aria-hidden="true" tabindex="-1"></a></span>
<span id="cb46-4"><a href="#cb46-4" aria-hidden="true" tabindex="-1"></a><span class="fu">use_swanlab</span><span class="kw">:</span><span class="at"> </span><span class="ch">true</span></span>
<span id="cb46-5"><a href="#cb46-5" aria-hidden="true" tabindex="-1"></a><span class="fu">swanlab_project</span><span class="kw">:</span><span class="at"> my-llm-project</span></span>
<span id="cb46-6"><a href="#cb46-6" aria-hidden="true" tabindex="-1"></a><span class="fu">swanlab_experiment_name</span><span class="kw">:</span><span class="at"> qwen-finetune-v1</span></span>
<span id="cb46-7"><a href="#cb46-7" aria-hidden="true" tabindex="-1"></a><span class="fu">swanlab_mode</span><span class="kw">:</span><span class="at"> cloud</span><span class="co"> # Options: cloud, local, offline, disabled</span></span>
<span id="cb46-8"><a href="#cb46-8" aria-hidden="true" tabindex="-1"></a><span class="fu">swanlab_workspace</span><span class="kw">:</span><span class="at"> my-team</span><span class="co"> # Optional: organization name</span></span>
<span id="cb46-9"><a href="#cb46-9" aria-hidden="true" tabindex="-1"></a><span class="fu">swanlab_api_key</span><span class="kw">:</span><span class="at"> YOUR_API_KEY</span><span class="co"> # Optional: can also use env var SWANLAB_API_KEY</span></span></code></pre></div><button title="Copy to Clipboard" class="code-copy-button"><i class="bi"></i></button></div>
</section>
<section id="run-training" class="level3">
<h3 class="anchored" data-anchor-id="run-training">3. Run Training</h3>
<div class="code-copy-outer-scaffold"><div class="sourceCode" id="cb47"><pre class="sourceCode bash code-with-copy"><code class="sourceCode bash"><span id="cb47-1"><a href="#cb47-1" aria-hidden="true" tabindex="-1"></a><span class="bu">export</span> <span class="va">SWANLAB_API_KEY</span><span class="op">=</span>your-api-key-here</span>
<span id="cb47-2"><a href="#cb47-2" aria-hidden="true" tabindex="-1"></a></span>
<span id="cb47-3"><a href="#cb47-3" aria-hidden="true" tabindex="-1"></a><span class="ex">swanlab</span> login</span>
<span id="cb47-4"><a href="#cb47-4" aria-hidden="true" tabindex="-1"></a></span>
<span id="cb47-5"><a href="#cb47-5" aria-hidden="true" tabindex="-1"></a><span class="ex">accelerate</span> launch <span class="at">-m</span> axolotl.cli.train your-config.yaml</span></code></pre></div><button title="Copy to Clipboard" class="code-copy-button"><i class="bi"></i></button></div>
</section>
<section id="configuration-options" class="level3">
<h3 class="anchored" data-anchor-id="configuration-options">Configuration Options</h3>
</section>
<section id="basic-configuration-1" class="level3">
<h3 class="anchored" data-anchor-id="basic-configuration-1">Basic Configuration</h3>
<table class="caption-top table">
<colgroup>
<col style="width: 28%">
<col style="width: 15%">
<col style="width: 23%">
<col style="width: 33%">
</colgroup>
<thead>
<tr class="header">
<th>Parameter</th>
<th>Type</th>
<th>Default</th>
<th>Description</th>
</tr>
</thead>
<tbody>
<tr class="odd">
<td><code>use_swanlab</code></td>
<td>bool</td>
<td><code>false</code></td>
<td>Enable SwanLab tracking</td>
</tr>
<tr class="even">
<td><code>swanlab_project</code></td>
<td>str</td>
<td><code>None</code></td>
<td>Project name (required)</td>
</tr>
<tr class="odd">
<td><code>swanlab_experiment_name</code></td>
<td>str</td>
<td><code>None</code></td>
<td>Experiment name</td>
</tr>
<tr class="even">
<td><code>swanlab_description</code></td>
<td>str</td>
<td><code>None</code></td>
<td>Experiment description</td>
</tr>
<tr class="odd">
<td><code>swanlab_mode</code></td>
<td>str</td>
<td><code>cloud</code></td>
<td>Sync mode: <code>cloud</code>, <code>local</code>, <code>offline</code>, <code>disabled</code></td>
</tr>
</tbody>
</table>
</section>
<section id="advanced-configuration" class="level3">
<h3 class="anchored" data-anchor-id="advanced-configuration">Advanced Configuration</h3>
<table class="caption-top table">
<colgroup>
<col style="width: 28%">
<col style="width: 15%">
<col style="width: 23%">
<col style="width: 33%">
</colgroup>
<thead>
<tr class="header">
<th>Parameter</th>
<th>Type</th>
<th>Default</th>
<th>Description</th>
</tr>
</thead>
<tbody>
<tr class="odd">
<td><code>swanlab_workspace</code></td>
<td>str</td>
<td><code>None</code></td>
<td>Workspace/organization name</td>
</tr>
<tr class="even">
<td><code>swanlab_api_key</code></td>
<td>str</td>
<td><code>None</code></td>
<td>API key (prefer env var)</td>
</tr>
<tr class="odd">
<td><code>swanlab_web_host</code></td>
<td>str</td>
<td><code>None</code></td>
<td>Private deployment web host</td>
</tr>
<tr class="even">
<td><code>swanlab_api_host</code></td>
<td>str</td>
<td><code>None</code></td>
<td>Private deployment API host</td>
</tr>
<tr class="odd">
<td><code>swanlab_log_model</code></td>
<td>bool</td>
<td><code>false</code></td>
<td>Log model checkpoints (coming soon)</td>
</tr>
<tr class="even">
<td><code>swanlab_lark_webhook_url</code></td>
<td>str</td>
<td><code>None</code></td>
<td>Lark (Feishu) webhook URL for team notifications</td>
</tr>
<tr class="odd">
<td><code>swanlab_lark_secret</code></td>
<td>str</td>
<td><code>None</code></td>
<td>Lark webhook HMAC secret for authentication</td>
</tr>
<tr class="even">
<td><code>swanlab_log_completions</code></td>
<td>bool</td>
<td><code>true</code></td>
<td>Enable RLHF completion table logging (DPO/KTO/ORPO/GRPO)</td>
</tr>
<tr class="odd">
<td><code>swanlab_completion_log_interval</code></td>
<td>int</td>
<td><code>100</code></td>
<td>Steps between completion logging</td>
</tr>
<tr class="even">
<td><code>swanlab_completion_max_buffer</code></td>
<td>int</td>
<td><code>128</code></td>
<td>Max completions to buffer (memory bound)</td>
</tr>
</tbody>
</table>
</section>
<section id="configuration-examples" class="level3">
<h3 class="anchored" data-anchor-id="configuration-examples">Configuration Examples</h3>
</section>
<section id="example-1-basic-cloud-sync" class="level3">
<h3 class="anchored" data-anchor-id="example-1-basic-cloud-sync">Example 1: Basic Cloud Sync</h3>
<div class="code-copy-outer-scaffold"><div class="sourceCode" id="cb48"><pre class="sourceCode yaml code-with-copy"><code class="sourceCode yaml"><span id="cb48-1"><a href="#cb48-1" aria-hidden="true" tabindex="-1"></a><span class="fu">plugins</span><span class="kw">:</span></span>
<span id="cb48-2"><a href="#cb48-2" aria-hidden="true" tabindex="-1"></a><span class="at"> </span><span class="kw">-</span><span class="at"> axolotl.integrations.swanlab.SwanLabPlugin</span></span>
<span id="cb48-3"><a href="#cb48-3" aria-hidden="true" tabindex="-1"></a></span>
<span id="cb48-4"><a href="#cb48-4" aria-hidden="true" tabindex="-1"></a><span class="fu">use_swanlab</span><span class="kw">:</span><span class="at"> </span><span class="ch">true</span></span>
<span id="cb48-5"><a href="#cb48-5" aria-hidden="true" tabindex="-1"></a><span class="fu">swanlab_project</span><span class="kw">:</span><span class="at"> llama-finetune</span></span>
<span id="cb48-6"><a href="#cb48-6" aria-hidden="true" tabindex="-1"></a><span class="fu">swanlab_experiment_name</span><span class="kw">:</span><span class="at"> llama-3-8b-instruct-v1</span></span>
<span id="cb48-7"><a href="#cb48-7" aria-hidden="true" tabindex="-1"></a><span class="fu">swanlab_mode</span><span class="kw">:</span><span class="at"> cloud</span></span></code></pre></div><button title="Copy to Clipboard" class="code-copy-button"><i class="bi"></i></button></div>
</section>
<section id="example-2-offlinelocal-mode" class="level3">
<h3 class="anchored" data-anchor-id="example-2-offlinelocal-mode">Example 2: Offline/Local Mode</h3>
<div class="code-copy-outer-scaffold"><div class="sourceCode" id="cb49"><pre class="sourceCode yaml code-with-copy"><code class="sourceCode yaml"><span id="cb49-1"><a href="#cb49-1" aria-hidden="true" tabindex="-1"></a><span class="fu">plugins</span><span class="kw">:</span></span>
<span id="cb49-2"><a href="#cb49-2" aria-hidden="true" tabindex="-1"></a><span class="at"> </span><span class="kw">-</span><span class="at"> axolotl.integrations.swanlab.SwanLabPlugin</span></span>
<span id="cb49-3"><a href="#cb49-3" aria-hidden="true" tabindex="-1"></a></span>
<span id="cb49-4"><a href="#cb49-4" aria-hidden="true" tabindex="-1"></a><span class="fu">use_swanlab</span><span class="kw">:</span><span class="at"> </span><span class="ch">true</span></span>
<span id="cb49-5"><a href="#cb49-5" aria-hidden="true" tabindex="-1"></a><span class="fu">swanlab_project</span><span class="kw">:</span><span class="at"> local-experiments</span></span>
<span id="cb49-6"><a href="#cb49-6" aria-hidden="true" tabindex="-1"></a><span class="fu">swanlab_experiment_name</span><span class="kw">:</span><span class="at"> test-run-1</span></span>
<span id="cb49-7"><a href="#cb49-7" aria-hidden="true" tabindex="-1"></a><span class="fu">swanlab_mode</span><span class="kw">:</span><span class="at"> local</span><span class="co"> # or 'offline'</span></span></code></pre></div><button title="Copy to Clipboard" class="code-copy-button"><i class="bi"></i></button></div>
</section>
<section id="example-3-team-workspace" class="level3">
<h3 class="anchored" data-anchor-id="example-3-team-workspace">Example 3: Team Workspace</h3>
<div class="code-copy-outer-scaffold"><div class="sourceCode" id="cb50"><pre class="sourceCode yaml code-with-copy"><code class="sourceCode yaml"><span id="cb50-1"><a href="#cb50-1" aria-hidden="true" tabindex="-1"></a><span class="fu">plugins</span><span class="kw">:</span></span>
<span id="cb50-2"><a href="#cb50-2" aria-hidden="true" tabindex="-1"></a><span class="at"> </span><span class="kw">-</span><span class="at"> axolotl.integrations.swanlab.SwanLabPlugin</span></span>
<span id="cb50-3"><a href="#cb50-3" aria-hidden="true" tabindex="-1"></a></span>
<span id="cb50-4"><a href="#cb50-4" aria-hidden="true" tabindex="-1"></a><span class="fu">use_swanlab</span><span class="kw">:</span><span class="at"> </span><span class="ch">true</span></span>
<span id="cb50-5"><a href="#cb50-5" aria-hidden="true" tabindex="-1"></a><span class="fu">swanlab_project</span><span class="kw">:</span><span class="at"> research-project</span></span>
<span id="cb50-6"><a href="#cb50-6" aria-hidden="true" tabindex="-1"></a><span class="fu">swanlab_experiment_name</span><span class="kw">:</span><span class="at"> experiment-42</span></span>
<span id="cb50-7"><a href="#cb50-7" aria-hidden="true" tabindex="-1"></a><span class="fu">swanlab_workspace</span><span class="kw">:</span><span class="at"> my-research-team</span></span>
<span id="cb50-8"><a href="#cb50-8" aria-hidden="true" tabindex="-1"></a><span class="fu">swanlab_mode</span><span class="kw">:</span><span class="at"> cloud</span></span></code></pre></div><button title="Copy to Clipboard" class="code-copy-button"><i class="bi"></i></button></div>
</section>
<section id="example-4-private-deployment" class="level3">
<h3 class="anchored" data-anchor-id="example-4-private-deployment">Example 4: Private Deployment</h3>
<div class="code-copy-outer-scaffold"><div class="sourceCode" id="cb51"><pre class="sourceCode yaml code-with-copy"><code class="sourceCode yaml"><span id="cb51-1"><a href="#cb51-1" aria-hidden="true" tabindex="-1"></a><span class="fu">plugins</span><span class="kw">:</span></span>
<span id="cb51-2"><a href="#cb51-2" aria-hidden="true" tabindex="-1"></a><span class="at"> </span><span class="kw">-</span><span class="at"> axolotl.integrations.swanlab.SwanLabPlugin</span></span>
<span id="cb51-3"><a href="#cb51-3" aria-hidden="true" tabindex="-1"></a></span>
<span id="cb51-4"><a href="#cb51-4" aria-hidden="true" tabindex="-1"></a><span class="fu">use_swanlab</span><span class="kw">:</span><span class="at"> </span><span class="ch">true</span></span>
<span id="cb51-5"><a href="#cb51-5" aria-hidden="true" tabindex="-1"></a><span class="fu">swanlab_project</span><span class="kw">:</span><span class="at"> internal-project</span></span>
<span id="cb51-6"><a href="#cb51-6" aria-hidden="true" tabindex="-1"></a><span class="fu">swanlab_experiment_name</span><span class="kw">:</span><span class="at"> secure-training</span></span>
<span id="cb51-7"><a href="#cb51-7" aria-hidden="true" tabindex="-1"></a><span class="fu">swanlab_mode</span><span class="kw">:</span><span class="at"> cloud</span></span>
<span id="cb51-8"><a href="#cb51-8" aria-hidden="true" tabindex="-1"></a><span class="fu">swanlab_web_host</span><span class="kw">:</span><span class="at"> https://swanlab.yourcompany.com</span></span>
<span id="cb51-9"><a href="#cb51-9" aria-hidden="true" tabindex="-1"></a><span class="fu">swanlab_api_host</span><span class="kw">:</span><span class="at"> https://api.swanlab.yourcompany.com</span></span></code></pre></div><button title="Copy to Clipboard" class="code-copy-button"><i class="bi"></i></button></div>
</section>
<section id="team-notifications-with-lark-feishu" class="level3">
<h3 class="anchored" data-anchor-id="team-notifications-with-lark-feishu">Team Notifications with Lark (Feishu)</h3>
<p>SwanLab supports sending real-time training notifications to your team chat via Lark (Feishu), ByteDance’s enterprise collaboration platform. This is especially useful for:
- <strong>Production training monitoring</strong>: Get alerts when training starts, completes, or encounters errors
- <strong>Team collaboration</strong>: Keep your ML team informed about long-running experiments
- <strong>Multi-timezone teams</strong>: Team members can check training progress without being online</p>
</section>
<section id="prerequisites-1" class="level3">
<h3 class="anchored" data-anchor-id="prerequisites-1">Prerequisites</h3>
<ol type="1">
<li><strong>Lark Bot Setup</strong>: Create a custom bot in your Lark group chat</li>
<li><strong>Webhook URL</strong>: Get the webhook URL from your Lark bot settings</li>
<li><strong>HMAC Secret</strong> (recommended): Enable signature verification in your Lark bot for security</li>
</ol>
<p>For detailed Lark bot setup instructions, see <a href="https://open.feishu.cn/document/ukTMukTMukTM/ucTM5YjL3ETO24yNxkjN">Lark Custom Bot Documentation</a>.</p>
</section>
<section id="example-5-basic-lark-notifications" class="level3">
<h3 class="anchored" data-anchor-id="example-5-basic-lark-notifications">Example 5: Basic Lark Notifications</h3>
<p>Send training notifications to a Lark group chat:</p>
<div class="code-copy-outer-scaffold"><div class="sourceCode" id="cb52"><pre class="sourceCode yaml code-with-copy"><code class="sourceCode yaml"><span id="cb52-1"><a href="#cb52-1" aria-hidden="true" tabindex="-1"></a><span class="fu">plugins</span><span class="kw">:</span></span>
<span id="cb52-2"><a href="#cb52-2" aria-hidden="true" tabindex="-1"></a><span class="at"> </span><span class="kw">-</span><span class="at"> axolotl.integrations.swanlab.SwanLabPlugin</span></span>
<span id="cb52-3"><a href="#cb52-3" aria-hidden="true" tabindex="-1"></a></span>
<span id="cb52-4"><a href="#cb52-4" aria-hidden="true" tabindex="-1"></a><span class="fu">use_swanlab</span><span class="kw">:</span><span class="at"> </span><span class="ch">true</span></span>
<span id="cb52-5"><a href="#cb52-5" aria-hidden="true" tabindex="-1"></a><span class="fu">swanlab_project</span><span class="kw">:</span><span class="at"> production-training</span></span>
<span id="cb52-6"><a href="#cb52-6" aria-hidden="true" tabindex="-1"></a><span class="fu">swanlab_experiment_name</span><span class="kw">:</span><span class="at"> llama-3-finetune-v2</span></span>
<span id="cb52-7"><a href="#cb52-7" aria-hidden="true" tabindex="-1"></a><span class="fu">swanlab_mode</span><span class="kw">:</span><span class="at"> cloud</span></span>
<span id="cb52-8"><a href="#cb52-8" aria-hidden="true" tabindex="-1"></a></span>
<span id="cb52-9"><a href="#cb52-9" aria-hidden="true" tabindex="-1"></a><span class="fu">swanlab_lark_webhook_url</span><span class="kw">:</span><span class="at"> https://open.feishu.cn/open-apis/bot/v2/hook/xxxxxxxxxx</span></span></code></pre></div><button title="Copy to Clipboard" class="code-copy-button"><i class="bi"></i></button></div>
<p><strong>Note</strong>: This configuration will work, but you’ll see a security warning recommending HMAC secret configuration.</p>
</section>
<section id="example-6-lark-notifications-with-hmac-security-recommended" class="level3">
<h3 class="anchored" data-anchor-id="example-6-lark-notifications-with-hmac-security-recommended">Example 6: Lark Notifications with HMAC Security (Recommended)</h3>
<p>For production use, enable HMAC signature verification:</p>
<div class="code-copy-outer-scaffold"><div class="sourceCode" id="cb53"><pre class="sourceCode yaml code-with-copy"><code class="sourceCode yaml"><span id="cb53-1"><a href="#cb53-1" aria-hidden="true" tabindex="-1"></a><span class="fu">plugins</span><span class="kw">:</span></span>
<span id="cb53-2"><a href="#cb53-2" aria-hidden="true" tabindex="-1"></a><span class="at"> </span><span class="kw">-</span><span class="at"> axolotl.integrations.swanlab.SwanLabPlugin</span></span>
<span id="cb53-3"><a href="#cb53-3" aria-hidden="true" tabindex="-1"></a></span>
<span id="cb53-4"><a href="#cb53-4" aria-hidden="true" tabindex="-1"></a><span class="fu">use_swanlab</span><span class="kw">:</span><span class="at"> </span><span class="ch">true</span></span>
<span id="cb53-5"><a href="#cb53-5" aria-hidden="true" tabindex="-1"></a><span class="fu">swanlab_project</span><span class="kw">:</span><span class="at"> production-training</span></span>
<span id="cb53-6"><a href="#cb53-6" aria-hidden="true" tabindex="-1"></a><span class="fu">swanlab_experiment_name</span><span class="kw">:</span><span class="at"> llama-3-finetune-v2</span></span>
<span id="cb53-7"><a href="#cb53-7" aria-hidden="true" tabindex="-1"></a><span class="fu">swanlab_mode</span><span class="kw">:</span><span class="at"> cloud</span></span>
<span id="cb53-8"><a href="#cb53-8" aria-hidden="true" tabindex="-1"></a></span>
<span id="cb53-9"><a href="#cb53-9" aria-hidden="true" tabindex="-1"></a><span class="fu">swanlab_lark_webhook_url</span><span class="kw">:</span><span class="at"> https://open.feishu.cn/open-apis/bot/v2/hook/xxxxxxxxxx</span></span>
<span id="cb53-10"><a href="#cb53-10" aria-hidden="true" tabindex="-1"></a><span class="fu">swanlab_lark_secret</span><span class="kw">:</span><span class="at"> your-webhook-secret-key</span></span></code></pre></div><button title="Copy to Clipboard" class="code-copy-button"><i class="bi"></i></button></div>
<p><strong>Why HMAC secret matters</strong>:
- Prevents unauthorized parties from sending fake notifications to your Lark group
- Ensures notifications genuinely come from your training jobs
- Required for production deployments with sensitive training data</p>
</section>
<section id="example-7-team-workspace-lark-notifications" class="level3">
<h3 class="anchored" data-anchor-id="example-7-team-workspace-lark-notifications">Example 7: Team Workspace + Lark Notifications</h3>
<p>Combine team workspace collaboration with Lark notifications:</p>
<div class="code-copy-outer-scaffold"><div class="sourceCode" id="cb54"><pre class="sourceCode yaml code-with-copy"><code class="sourceCode yaml"><span id="cb54-1"><a href="#cb54-1" aria-hidden="true" tabindex="-1"></a><span class="fu">plugins</span><span class="kw">:</span></span>
<span id="cb54-2"><a href="#cb54-2" aria-hidden="true" tabindex="-1"></a><span class="at"> </span><span class="kw">-</span><span class="at"> axolotl.integrations.swanlab.SwanLabPlugin</span></span>
<span id="cb54-3"><a href="#cb54-3" aria-hidden="true" tabindex="-1"></a></span>
<span id="cb54-4"><a href="#cb54-4" aria-hidden="true" tabindex="-1"></a><span class="fu">use_swanlab</span><span class="kw">:</span><span class="at"> </span><span class="ch">true</span></span>
<span id="cb54-5"><a href="#cb54-5" aria-hidden="true" tabindex="-1"></a><span class="fu">swanlab_project</span><span class="kw">:</span><span class="at"> research-project</span></span>
<span id="cb54-6"><a href="#cb54-6" aria-hidden="true" tabindex="-1"></a><span class="fu">swanlab_experiment_name</span><span class="kw">:</span><span class="at"> multimodal-experiment-42</span></span>
<span id="cb54-7"><a href="#cb54-7" aria-hidden="true" tabindex="-1"></a><span class="fu">swanlab_workspace</span><span class="kw">:</span><span class="at"> ml-research-team</span></span>
<span id="cb54-8"><a href="#cb54-8" aria-hidden="true" tabindex="-1"></a><span class="fu">swanlab_mode</span><span class="kw">:</span><span class="at"> cloud</span></span>
<span id="cb54-9"><a href="#cb54-9" aria-hidden="true" tabindex="-1"></a></span>
<span id="cb54-10"><a href="#cb54-10" aria-hidden="true" tabindex="-1"></a><span class="fu">swanlab_lark_webhook_url</span><span class="kw">:</span><span class="at"> https://open.feishu.cn/open-apis/bot/v2/hook/xxxxxxxxxx</span></span>
<span id="cb54-11"><a href="#cb54-11" aria-hidden="true" tabindex="-1"></a><span class="fu">swanlab_lark_secret</span><span class="kw">:</span><span class="at"> your-webhook-secret-key</span></span></code></pre></div><button title="Copy to Clipboard" class="code-copy-button"><i class="bi"></i></button></div>
</section>
<section id="what-notifications-are-sent" class="level3">
<h3 class="anchored" data-anchor-id="what-notifications-are-sent">What Notifications Are Sent?</h3>
<p>SwanLab’s Lark integration sends notifications for key training events:
- <strong>Training Start</strong>: When your experiment begins
- <strong>Training Complete</strong>: When training finishes successfully
- <strong>Training Errors</strong>: If training crashes or encounters critical errors
- <strong>Metric Milestones</strong>: Configurable alerts for metric thresholds (if configured in SwanLab)</p>
<p>Each notification includes:
- Experiment name and project
- Training status
- Key metrics (loss, learning rate)
- Direct link to SwanLab dashboard</p>
</section>
<section id="lark-configuration-validation" class="level3">
<h3 class="anchored" data-anchor-id="lark-configuration-validation">Lark Configuration Validation</h3>
<p>The plugin validates your Lark configuration at startup:</p>
<section id="valid-configurations" class="level4">
<h4 class="anchored" data-anchor-id="valid-configurations">✅ Valid Configurations</h4>
<div class="code-copy-outer-scaffold"><div class="sourceCode" id="cb55"><pre class="sourceCode yaml code-with-copy"><code class="sourceCode yaml"><span id="cb55-1"><a href="#cb55-1" aria-hidden="true" tabindex="-1"></a><span class="fu">use_swanlab</span><span class="kw">:</span><span class="at"> </span><span class="ch">true</span></span>
<span id="cb55-2"><a href="#cb55-2" aria-hidden="true" tabindex="-1"></a><span class="fu">swanlab_project</span><span class="kw">:</span><span class="at"> my-project</span></span>
<span id="cb55-3"><a href="#cb55-3" aria-hidden="true" tabindex="-1"></a></span>
<span id="cb55-4"><a href="#cb55-4" aria-hidden="true" tabindex="-1"></a><span class="fu">use_swanlab</span><span class="kw">:</span><span class="at"> </span><span class="ch">true</span></span>
<span id="cb55-5"><a href="#cb55-5" aria-hidden="true" tabindex="-1"></a><span class="fu">swanlab_project</span><span class="kw">:</span><span class="at"> my-project</span></span>
<span id="cb55-6"><a href="#cb55-6" aria-hidden="true" tabindex="-1"></a><span class="fu">swanlab_lark_webhook_url</span><span class="kw">:</span><span class="at"> https://open.feishu.cn/open-apis/bot/v2/hook/xxx</span></span>
<span id="cb55-7"><a href="#cb55-7" aria-hidden="true" tabindex="-1"></a><span class="fu">swanlab_lark_secret</span><span class="kw">:</span><span class="at"> your-secret</span></span>
<span id="cb55-8"><a href="#cb55-8" aria-hidden="true" tabindex="-1"></a></span>
<span id="cb55-9"><a href="#cb55-9" aria-hidden="true" tabindex="-1"></a><span class="fu">use_swanlab</span><span class="kw">:</span><span class="at"> </span><span class="ch">true</span></span>
<span id="cb55-10"><a href="#cb55-10" aria-hidden="true" tabindex="-1"></a><span class="fu">swanlab_project</span><span class="kw">:</span><span class="at"> my-project</span></span>
<span id="cb55-11"><a href="#cb55-11" aria-hidden="true" tabindex="-1"></a><span class="fu">swanlab_lark_webhook_url</span><span class="kw">:</span><span class="at"> https://open.feishu.cn/open-apis/bot/v2/hook/xxx</span></span></code></pre></div><button title="Copy to Clipboard" class="code-copy-button"><i class="bi"></i></button></div>
</section>
</section>
<section id="security-best-practices" class="level3">
<h3 class="anchored" data-anchor-id="security-best-practices">Security Best Practices</h3>
<ol type="1">
<li><p><strong>Always use HMAC secret in production</strong>:</p>
<div class="code-copy-outer-scaffold"><div class="sourceCode" id="cb56"><pre class="sourceCode yaml code-with-copy"><code class="sourceCode yaml"><span id="cb56-1"><a href="#cb56-1" aria-hidden="true" tabindex="-1"></a><span class="fu">swanlab_lark_webhook_url</span><span class="kw">:</span><span class="at"> https://open.feishu.cn/...</span></span>
<span id="cb56-2"><a href="#cb56-2" aria-hidden="true" tabindex="-1"></a><span class="fu">swanlab_lark_secret</span><span class="kw">:</span><span class="at"> your-secret-key</span><span class="co"> # ✅ Add this!</span></span></code></pre></div><button title="Copy to Clipboard" class="code-copy-button"><i class="bi"></i></button></div></li>
<li><p><strong>Store secrets in environment variables</strong> (even better):</p>
<div class="code-copy-outer-scaffold"><div class="sourceCode" id="cb57"><pre class="sourceCode yaml code-with-copy"><code class="sourceCode yaml"><span id="cb57-1"><a href="#cb57-1" aria-hidden="true" tabindex="-1"></a><span class="co"># In your training script/environment</span></span>
<span id="cb57-2"><a href="#cb57-2" aria-hidden="true" tabindex="-1"></a><span class="at">export SWANLAB_LARK_WEBHOOK_URL="https://open.feishu.cn/..."</span></span>
<span id="cb57-3"><a href="#cb57-3" aria-hidden="true" tabindex="-1"></a><span class="at">export SWANLAB_LARK_SECRET="your-secret-key"</span></span></code></pre></div><button title="Copy to Clipboard" class="code-copy-button"><i class="bi"></i></button></div>
<p>Then in config:</p>
<div class="code-copy-outer-scaffold"><div class="sourceCode" id="cb58"><pre class="sourceCode yaml code-with-copy"><code class="sourceCode yaml"><span id="cb58-1"><a href="#cb58-1" aria-hidden="true" tabindex="-1"></a><span class="co"># SwanLab plugin will auto-detect environment variables</span></span>
<span id="cb58-2"><a href="#cb58-2" aria-hidden="true" tabindex="-1"></a><span class="fu">use_swanlab</span><span class="kw">:</span><span class="at"> </span><span class="ch">true</span></span>
<span id="cb58-3"><a href="#cb58-3" aria-hidden="true" tabindex="-1"></a><span class="fu">swanlab_project</span><span class="kw">:</span><span class="at"> my-project</span></span>
<span id="cb58-4"><a href="#cb58-4" aria-hidden="true" tabindex="-1"></a><span class="co"># Lark URL and secret read from env vars</span></span></code></pre></div><button title="Copy to Clipboard" class="code-copy-button"><i class="bi"></i></button></div></li>
<li><p><strong>Rotate webhook secrets periodically</strong>: Update your Lark bot’s secret every 90 days</p></li>
<li><p><strong>Use separate webhooks for dev/prod</strong>: Don’t mix development and production notifications</p></li>
</ol>
</section>
<section id="distributed-training" class="level3">
<h3 class="anchored" data-anchor-id="distributed-training">Distributed Training</h3>
<p>Lark notifications are automatically deduplicated in distributed training:
- Only <strong>rank 0</strong> sends notifications
- Other GPU ranks skip Lark registration
- Prevents duplicate messages in multi-GPU training</p>
<div class="code-copy-outer-scaffold"><div class="sourceCode" id="cb59"><pre class="sourceCode bash code-with-copy"><code class="sourceCode bash"><span id="cb59-1"><a href="#cb59-1" aria-hidden="true" tabindex="-1"></a><span class="ex">torchrun</span> <span class="at">--nproc_per_node</span><span class="op">=</span>4 <span class="at">-m</span> axolotl.cli.train config.yml</span></code></pre></div><button title="Copy to Clipboard" class="code-copy-button"><i class="bi"></i></button></div>
</section>
<section id="rlhf-completion-table-logging" class="level3">
<h3 class="anchored" data-anchor-id="rlhf-completion-table-logging">RLHF Completion Table Logging</h3>
<p>For RLHF (Reinforcement Learning from Human Feedback) training methods like DPO, KTO, ORPO, and GRPO, SwanLab can log model completions (prompts, chosen/rejected responses, rewards) to a visual table for qualitative analysis. This helps you:</p>
<ul>
<li><strong>Inspect model behavior</strong>: See actual model outputs during training</li>
<li><strong>Debug preference learning</strong>: Compare chosen vs rejected responses</li>
<li><strong>Track reward patterns</strong>: Monitor how rewards evolve over training</li>
<li><strong>Share examples with team</strong>: Visual tables in SwanLab dashboard</li>
</ul>
</section>
<section id="features-1" class="level3">
<h3 class="anchored" data-anchor-id="features-1">Features</h3>
<ul>
<li>✅ <strong>Automatic detection</strong>: Works with DPO, KTO, ORPO, GRPO trainers</li>
<li>✅ <strong>Memory-safe buffering</strong>: Bounded buffer prevents memory leaks in long training runs</li>
<li>✅ <strong>Periodic logging</strong>: Configurable logging interval to reduce overhead</li>
<li>✅ <strong>Rich visualization</strong>: SwanLab tables show prompts, responses, and metrics side-by-side</li>
</ul>
</section>
<section id="configuration" class="level3">
<h3 class="anchored" data-anchor-id="configuration">Configuration</h3>
<table class="caption-top table">
<colgroup>
<col style="width: 28%">
<col style="width: 15%">
<col style="width: 23%">
<col style="width: 33%">
</colgroup>
<thead>
<tr class="header">
<th>Parameter</th>
<th>Type</th>
<th>Default</th>
<th>Description</th>
</tr>
</thead>
<tbody>
<tr class="odd">
<td><code>swanlab_log_completions</code></td>
<td>bool</td>
<td><code>true</code></td>
<td>Enable completion logging for RLHF trainers</td>
</tr>
<tr class="even">
<td><code>swanlab_completion_log_interval</code></td>
<td>int</td>
<td><code>100</code></td>
<td>Log completions to SwanLab every N training steps</td>
</tr>
<tr class="odd">
<td><code>swanlab_completion_max_buffer</code></td>
<td>int</td>
<td><code>128</code></td>
<td>Maximum completions to buffer (memory bound)</td>
</tr>
</tbody>
</table>
</section>
<section id="example-dpo-training-with-completion-logging" class="level3">
<h3 class="anchored" data-anchor-id="example-dpo-training-with-completion-logging">Example: DPO Training with Completion Logging</h3>
<div class="code-copy-outer-scaffold"><div class="sourceCode" id="cb60"><pre class="sourceCode yaml code-with-copy"><code class="sourceCode yaml"><span id="cb60-1"><a href="#cb60-1" aria-hidden="true" tabindex="-1"></a><span class="fu">plugins</span><span class="kw">:</span></span>
<span id="cb60-2"><a href="#cb60-2" aria-hidden="true" tabindex="-1"></a><span class="at"> </span><span class="kw">-</span><span class="at"> axolotl.integrations.swanlab.SwanLabPlugin</span></span>
<span id="cb60-3"><a href="#cb60-3" aria-hidden="true" tabindex="-1"></a></span>
<span id="cb60-4"><a href="#cb60-4" aria-hidden="true" tabindex="-1"></a><span class="fu">use_swanlab</span><span class="kw">:</span><span class="at"> </span><span class="ch">true</span></span>
<span id="cb60-5"><a href="#cb60-5" aria-hidden="true" tabindex="-1"></a><span class="fu">swanlab_project</span><span class="kw">:</span><span class="at"> dpo-training</span></span>
<span id="cb60-6"><a href="#cb60-6" aria-hidden="true" tabindex="-1"></a><span class="fu">swanlab_experiment_name</span><span class="kw">:</span><span class="at"> llama-3-dpo-v1</span></span>
<span id="cb60-7"><a href="#cb60-7" aria-hidden="true" tabindex="-1"></a><span class="fu">swanlab_mode</span><span class="kw">:</span><span class="at"> cloud</span></span>
<span id="cb60-8"><a href="#cb60-8" aria-hidden="true" tabindex="-1"></a></span>
<span id="cb60-9"><a href="#cb60-9" aria-hidden="true" tabindex="-1"></a><span class="fu">swanlab_log_completions</span><span class="kw">:</span><span class="at"> </span><span class="ch">true</span></span>
<span id="cb60-10"><a href="#cb60-10" aria-hidden="true" tabindex="-1"></a><span class="fu">swanlab_completion_log_interval</span><span class="kw">:</span><span class="at"> </span><span class="dv">100</span><span class="co"> # Log every 100 steps</span></span>
<span id="cb60-11"><a href="#cb60-11" aria-hidden="true" tabindex="-1"></a><span class="fu">swanlab_completion_max_buffer</span><span class="kw">:</span><span class="at"> </span><span class="dv">128</span><span class="co"> # Keep last 128 completions</span></span>
<span id="cb60-12"><a href="#cb60-12" aria-hidden="true" tabindex="-1"></a></span>
<span id="cb60-13"><a href="#cb60-13" aria-hidden="true" tabindex="-1"></a><span class="fu">rl</span><span class="kw">:</span><span class="at"> dpo</span></span>
<span id="cb60-14"><a href="#cb60-14" aria-hidden="true" tabindex="-1"></a><span class="fu">datasets</span><span class="kw">:</span></span>
<span id="cb60-15"><a href="#cb60-15" aria-hidden="true" tabindex="-1"></a><span class="at"> </span><span class="kw">-</span><span class="at"> </span><span class="fu">path</span><span class="kw">:</span><span class="at"> /path/to/preference_dataset</span></span>
<span id="cb60-16"><a href="#cb60-16" aria-hidden="true" tabindex="-1"></a><span class="at"> </span><span class="fu">type</span><span class="kw">:</span><span class="at"> chatml.intel</span></span></code></pre></div><button title="Copy to Clipboard" class="code-copy-button"><i class="bi"></i></button></div>
</section>
<section id="example-disable-completion-logging" class="level3">
<h3 class="anchored" data-anchor-id="example-disable-completion-logging">Example: Disable Completion Logging</h3>
<p>If you’re doing a quick test run or don’t need completion tables:</p>
<div class="code-copy-outer-scaffold"><div class="sourceCode" id="cb61"><pre class="sourceCode yaml code-with-copy"><code class="sourceCode yaml"><span id="cb61-1"><a href="#cb61-1" aria-hidden="true" tabindex="-1"></a><span class="fu">plugins</span><span class="kw">:</span></span>
<span id="cb61-2"><a href="#cb61-2" aria-hidden="true" tabindex="-1"></a><span class="at"> </span><span class="kw">-</span><span class="at"> axolotl.integrations.swanlab.SwanLabPlugin</span></span>
<span id="cb61-3"><a href="#cb61-3" aria-hidden="true" tabindex="-1"></a></span>
<span id="cb61-4"><a href="#cb61-4" aria-hidden="true" tabindex="-1"></a><span class="fu">use_swanlab</span><span class="kw">:</span><span class="at"> </span><span class="ch">true</span></span>
<span id="cb61-5"><a href="#cb61-5" aria-hidden="true" tabindex="-1"></a><span class="fu">swanlab_project</span><span class="kw">:</span><span class="at"> dpo-training</span></span>
<span id="cb61-6"><a href="#cb61-6" aria-hidden="true" tabindex="-1"></a></span>
<span id="cb61-7"><a href="#cb61-7" aria-hidden="true" tabindex="-1"></a><span class="fu">swanlab_log_completions</span><span class="kw">:</span><span class="at"> </span><span class="ch">false</span></span></code></pre></div><button title="Copy to Clipboard" class="code-copy-button"><i class="bi"></i></button></div>
</section>
<section id="supported-rlhf-trainers" class="level3">
<h3 class="anchored" data-anchor-id="supported-rlhf-trainers">Supported RLHF Trainers</h3>
<p>The completion logging callback automatically activates for these trainer types:</p>
<ul>
<li><strong>DPO (Direct Preference Optimization)</strong>: Logs prompts, chosen, rejected, reward_diff</li>
<li><strong>KTO (Kahneman-Tversky Optimization)</strong>: Logs prompts, completions, labels, rewards</li>
<li><strong>ORPO (Odds Ratio Preference Optimization)</strong>: Logs prompts, chosen, rejected, log_odds_ratio</li>
<li><strong>GRPO (Group Relative Policy Optimization)</strong>: Logs prompts, completions, rewards, advantages</li>
<li><strong>CPO (Constrained Policy Optimization)</strong>: Logs prompts, chosen, rejected</li>
</ul>
<p>For non-RLHF trainers (standard supervised fine-tuning), the completion callback is automatically skipped.</p>
</section>
<section id="how-it-works-3" class="level3">
<h3 class="anchored" data-anchor-id="how-it-works-3">How It Works</h3>
<ol type="1">
<li><strong>Auto-detection</strong>: Plugin detects trainer type at initialization</li>
<li><strong>Buffering</strong>: Completions are buffered in memory (up to <code>swanlab_completion_max_buffer</code>)</li>
<li><strong>Periodic logging</strong>: Every <code>swanlab_completion_log_interval</code> steps, buffer is logged to SwanLab</li>
<li><strong>Memory safety</strong>: Old completions are automatically dropped when buffer is full (uses <code>collections.deque</code>)</li>
<li><strong>Final flush</strong>: Remaining completions are logged when training completes</li>
</ol>
</section>
<section id="viewing-completion-tables" class="level3">
<h3 class="anchored" data-anchor-id="viewing-completion-tables">Viewing Completion Tables</h3>
<p>After training starts, you can view completion tables in your SwanLab dashboard:</p>
<ol type="1">
<li>Navigate to your experiment in SwanLab</li>
<li>Look for the “rlhf_completions” table in the metrics panel</li>
<li>The table shows:
<ul>
<li><strong>step</strong>: Training step when completion was generated</li>
<li><strong>prompt</strong>: Input prompt</li>
<li><strong>chosen</strong>: Preferred response (DPO/ORPO)</li>
<li><strong>rejected</strong>: Non-preferred response (DPO/ORPO)</li>
<li><strong>completion</strong>: Model output (KTO/GRPO)</li>
<li><strong>reward_diff/reward</strong>: Reward metrics</li>
<li>Trainer-specific metrics (e.g., log_odds_ratio for ORPO)</li>
</ul></li>
</ol>
</section>
<section id="memory-management" class="level3">
<h3 class="anchored" data-anchor-id="memory-management">Memory Management</h3>
<p>The completion buffer is <strong>memory-bounded</strong> to prevent memory leaks:</p>
<div class="code-copy-outer-scaffold"><div class="sourceCode" id="cb62"><pre class="sourceCode python code-with-copy"><code class="sourceCode python"><span id="cb62-1"><a href="#cb62-1" aria-hidden="true" tabindex="-1"></a><span class="im">from</span> collections <span class="im">import</span> deque</span>
<span id="cb62-2"><a href="#cb62-2" aria-hidden="true" tabindex="-1"></a></span>
<span id="cb62-3"><a href="#cb62-3" aria-hidden="true" tabindex="-1"></a><span class="bu">buffer</span> <span class="op">=</span> deque(maxlen<span class="op">=</span><span class="dv">128</span>) <span class="co"># Old completions automatically dropped</span></span></code></pre></div><button title="Copy to Clipboard" class="code-copy-button"><i class="bi"></i></button></div>
<p><strong>Memory usage estimate</strong>:
- Average completion: ~500 characters (prompt + responses)
- Buffer size 128: ~64 KB (negligible)
- Buffer size 1024: ~512 KB (still small)</p>
<p><strong>Recommendation</strong>: Default buffer size (128) works well for most cases. Increase to 512-1024 only if you need to review more historical completions.</p>
</section>
<section id="performance-impact" class="level3">
<h3 class="anchored" data-anchor-id="performance-impact">Performance Impact</h3>
<p>Completion logging has minimal overhead:</p>
<ul>
<li><strong>Buffering</strong>: O(1) append operation, negligible CPU/memory</li>
<li><strong>Logging</strong>: Only happens every N steps (default: 100)</li>
<li><strong>Network</strong>: SwanLab batches table uploads efficiently</li>
</ul>
<p><strong>Expected overhead</strong>: &lt; 0.5% per training step</p>
</section>
<section id="troubleshooting" class="level3">
<h3 class="anchored" data-anchor-id="troubleshooting">Troubleshooting</h3>
<section id="completions-not-appearing-in-swanlab" class="level4">
<h4 class="anchored" data-anchor-id="completions-not-appearing-in-swanlab">Completions not appearing in SwanLab</h4>
<p><strong>Cause</strong>: Trainer may not be logging completion data in the expected format.</p>
<p><strong>Diagnostic steps</strong>:
1. Check trainer type detection in logs:
<code>text INFO: SwanLab RLHF completion logging enabled for DPOTrainer (type: dpo)</code>
2. Verify your trainer is an RLHF trainer (DPO/KTO/ORPO/GRPO)
3. Check if trainer logs completion data (this depends on TRL version)</p>
<p><strong>Note</strong>: The current implementation expects trainers to log completion data in the <code>logs</code> dict during <code>on_log()</code> callback. Some TRL trainers may not expose this data by default. You may need to patch the trainer to expose completions.</p>
</section>
<section id="buffer-fills-up-too-quickly" class="level4">
<h4 class="anchored" data-anchor-id="buffer-fills-up-too-quickly">Buffer fills up too quickly</h4>
<p><strong>Cause</strong>: High logging frequency with small buffer size.</p>
<p><strong>Solution</strong>: Increase buffer size or logging interval:</p>
<div class="code-copy-outer-scaffold"><div class="sourceCode" id="cb63"><pre class="sourceCode yaml code-with-copy"><code class="sourceCode yaml"><span id="cb63-1"><a href="#cb63-1" aria-hidden="true" tabindex="-1"></a><span class="fu">swanlab_completion_log_interval</span><span class="kw">:</span><span class="at"> </span><span class="dv">200</span><span class="co"> # Log less frequently</span></span>
<span id="cb63-2"><a href="#cb63-2" aria-hidden="true" tabindex="-1"></a><span class="fu">swanlab_completion_max_buffer</span><span class="kw">:</span><span class="at"> </span><span class="dv">512</span><span class="co"> # Larger buffer</span></span></code></pre></div><button title="Copy to Clipboard" class="code-copy-button"><i class="bi"></i></button></div>
</section>
<section id="memory-usage-growing-over-time" class="level4">
<h4 class="anchored" data-anchor-id="memory-usage-growing-over-time">Memory usage growing over time</h4>
<p><strong>Cause</strong>: Buffer should be bounded, so this indicates a bug.</p>
<p><strong>Solution</strong>:
1. Verify <code>swanlab_completion_max_buffer</code> is set
2. Check SwanLab version is up to date
3. Report issue with memory profiling data</p>
</section>
</section>
<section id="performance-profiling" class="level3">
<h3 class="anchored" data-anchor-id="performance-profiling">Performance Profiling</h3>
<p>SwanLab integration includes profiling utilities to measure and log execution time of trainer methods. This helps you:</p>
<ul>
<li><strong>Identify bottlenecks</strong>: Find slow operations in your training loop</li>
<li><strong>Optimize performance</strong>: Track improvements after optimization changes</li>
<li><strong>Monitor distributed training</strong>: See per-rank timing differences</li>
<li><strong>Debug hangs</strong>: Detect methods that take unexpectedly long</li>
</ul>
</section>
<section id="features-2" class="level3">
<h3 class="anchored" data-anchor-id="features-2">Features</h3>
<ul>
<li>✅ <strong>Zero-config profiling</strong>: Automatic timing of key trainer methods</li>
<li>✅ <strong>Decorator-based</strong>: Easy to add profiling to custom methods with <code>@swanlab_profile</code></li>
<li>✅ <strong>Context manager</strong>: Fine-grained profiling with <code>swanlab_profiling_context()</code></li>
<li>✅ <strong>Advanced filtering</strong>: <code>ProfilingConfig</code> for throttling and minimum duration thresholds</li>
<li>✅ <strong>Exception-safe</strong>: Logs duration even if function raises an exception</li>
</ul>
</section>
<section id="basic-usage-decorator" class="level3">
<h3 class="anchored" data-anchor-id="basic-usage-decorator">Basic Usage: Decorator</h3>
<p>Add profiling to any trainer method with the <code>@swanlab_profile</code> decorator:</p>
<div class="code-copy-outer-scaffold"><div class="sourceCode" id="cb64"><pre class="sourceCode python code-with-copy"><code class="sourceCode python"><span id="cb64-1"><a href="#cb64-1" aria-hidden="true" tabindex="-1"></a><span class="im">from</span> axolotl.integrations.swanlab.profiling <span class="im">import</span> swanlab_profile</span>
<span id="cb64-2"><a href="#cb64-2" aria-hidden="true" tabindex="-1"></a></span>
<span id="cb64-3"><a href="#cb64-3" aria-hidden="true" tabindex="-1"></a><span class="kw">class</span> MyCustomTrainer(AxolotlTrainer):</span>
<span id="cb64-4"><a href="#cb64-4" aria-hidden="true" tabindex="-1"></a> <span class="at">@swanlab_profile</span></span>
<span id="cb64-5"><a href="#cb64-5" aria-hidden="true" tabindex="-1"></a> <span class="kw">def</span> training_step(<span class="va">self</span>, model, inputs):</span>
<span id="cb64-6"><a href="#cb64-6" aria-hidden="true" tabindex="-1"></a> <span class="co"># Your training step logic</span></span>
<span id="cb64-7"><a href="#cb64-7" aria-hidden="true" tabindex="-1"></a> <span class="cf">return</span> <span class="bu">super</span>().training_step(model, inputs)</span>
<span id="cb64-8"><a href="#cb64-8" aria-hidden="true" tabindex="-1"></a></span>
<span id="cb64-9"><a href="#cb64-9" aria-hidden="true" tabindex="-1"></a> <span class="at">@swanlab_profile</span></span>
<span id="cb64-10"><a href="#cb64-10" aria-hidden="true" tabindex="-1"></a> <span class="kw">def</span> prediction_step(<span class="va">self</span>, model, inputs, prediction_loss_only):</span>
<span id="cb64-11"><a href="#cb64-11" aria-hidden="true" tabindex="-1"></a> <span class="co"># Your prediction logic</span></span>
<span id="cb64-12"><a href="#cb64-12" aria-hidden="true" tabindex="-1"></a> <span class="cf">return</span> <span class="bu">super</span>().prediction_step(model, inputs, prediction_loss_only)</span></code></pre></div><button title="Copy to Clipboard" class="code-copy-button"><i class="bi"></i></button></div>
<p>The decorator automatically:
1. Measures execution time with high-precision timer
2. Logs to SwanLab as <code>profiling/Time taken: ClassName.method_name</code>
3. Only logs if SwanLab is enabled (<code>use_swanlab: true</code>)
4. Gracefully handles exceptions (logs duration, then re-raises)</p>
</section>
<section id="advanced-usage-context-manager" class="level3">
<h3 class="anchored" data-anchor-id="advanced-usage-context-manager">Advanced Usage: Context Manager</h3>
<p>For fine-grained profiling within a method:</p>
<div class="code-copy-outer-scaffold"><div class="sourceCode" id="cb65"><pre class="sourceCode python code-with-copy"><code class="sourceCode python"><span id="cb65-1"><a href="#cb65-1" aria-hidden="true" tabindex="-1"></a><span class="im">from</span> axolotl.integrations.swanlab.profiling <span class="im">import</span> swanlab_profiling_context</span>
<span id="cb65-2"><a href="#cb65-2" aria-hidden="true" tabindex="-1"></a></span>
<span id="cb65-3"><a href="#cb65-3" aria-hidden="true" tabindex="-1"></a><span class="kw">class</span> MyTrainer(AxolotlTrainer):</span>
<span id="cb65-4"><a href="#cb65-4" aria-hidden="true" tabindex="-1"></a> <span class="kw">def</span> complex_training_step(<span class="va">self</span>, model, inputs):</span>
<span id="cb65-5"><a href="#cb65-5" aria-hidden="true" tabindex="-1"></a> <span class="co"># Profile just the forward pass</span></span>
<span id="cb65-6"><a href="#cb65-6" aria-hidden="true" tabindex="-1"></a> <span class="cf">with</span> swanlab_profiling_context(<span class="va">self</span>, <span class="st">"forward_pass"</span>):</span>
<span id="cb65-7"><a href="#cb65-7" aria-hidden="true" tabindex="-1"></a> outputs <span class="op">=</span> model(<span class="op">**</span>inputs)</span>
<span id="cb65-8"><a href="#cb65-8" aria-hidden="true" tabindex="-1"></a></span>
<span id="cb65-9"><a href="#cb65-9" aria-hidden="true" tabindex="-1"></a> <span class="co"># Profile just the backward pass</span></span>
<span id="cb65-10"><a href="#cb65-10" aria-hidden="true" tabindex="-1"></a> <span class="cf">with</span> swanlab_profiling_context(<span class="va">self</span>, <span class="st">"backward_pass"</span>):</span>
<span id="cb65-11"><a href="#cb65-11" aria-hidden="true" tabindex="-1"></a> loss <span class="op">=</span> outputs.loss</span>
<span id="cb65-12"><a href="#cb65-12" aria-hidden="true" tabindex="-1"></a> loss.backward()</span>
<span id="cb65-13"><a href="#cb65-13" aria-hidden="true" tabindex="-1"></a></span>
<span id="cb65-14"><a href="#cb65-14" aria-hidden="true" tabindex="-1"></a> <span class="cf">return</span> outputs</span></code></pre></div><button title="Copy to Clipboard" class="code-copy-button"><i class="bi"></i></button></div>
</section>
<section id="advanced-usage-profilingconfig" class="level3">
<h3 class="anchored" data-anchor-id="advanced-usage-profilingconfig">Advanced Usage: ProfilingConfig</h3>
<p>Filter and throttle profiling logs with <code>ProfilingConfig</code>:</p>
<div class="code-copy-outer-scaffold"><div class="sourceCode" id="cb66"><pre class="sourceCode python code-with-copy"><code class="sourceCode python"><span id="cb66-1"><a href="#cb66-1" aria-hidden="true" tabindex="-1"></a><span class="im">from</span> axolotl.integrations.swanlab.profiling <span class="im">import</span> (</span>
<span id="cb66-2"><a href="#cb66-2" aria-hidden="true" tabindex="-1"></a> swanlab_profiling_context_advanced,</span>
<span id="cb66-3"><a href="#cb66-3" aria-hidden="true" tabindex="-1"></a> ProfilingConfig,</span>
<span id="cb66-4"><a href="#cb66-4" aria-hidden="true" tabindex="-1"></a>)</span>
<span id="cb66-5"><a href="#cb66-5" aria-hidden="true" tabindex="-1"></a></span>
<span id="cb66-6"><a href="#cb66-6" aria-hidden="true" tabindex="-1"></a>profiling_config <span class="op">=</span> ProfilingConfig(</span>
<span id="cb66-7"><a href="#cb66-7" aria-hidden="true" tabindex="-1"></a> enabled<span class="op">=</span><span class="va">True</span>,</span>
<span id="cb66-8"><a href="#cb66-8" aria-hidden="true" tabindex="-1"></a> min_duration_ms<span class="op">=</span><span class="fl">1.0</span>, <span class="co"># Only log if duration &gt; 1ms</span></span>
<span id="cb66-9"><a href="#cb66-9" aria-hidden="true" tabindex="-1"></a> log_interval<span class="op">=</span><span class="dv">10</span>, <span class="co"># Log every 10th call</span></span>
<span id="cb66-10"><a href="#cb66-10" aria-hidden="true" tabindex="-1"></a>)</span>
<span id="cb66-11"><a href="#cb66-11" aria-hidden="true" tabindex="-1"></a></span>
<span id="cb66-12"><a href="#cb66-12" aria-hidden="true" tabindex="-1"></a><span class="kw">class</span> MyTrainer(AxolotlTrainer):</span>
<span id="cb66-13"><a href="#cb66-13" aria-hidden="true" tabindex="-1"></a> <span class="kw">def</span> frequently_called_method(<span class="va">self</span>, data):</span>
<span id="cb66-14"><a href="#cb66-14" aria-hidden="true" tabindex="-1"></a> <span class="cf">with</span> swanlab_profiling_context_advanced(</span>
<span id="cb66-15"><a href="#cb66-15" aria-hidden="true" tabindex="-1"></a> <span class="va">self</span>,</span>
<span id="cb66-16"><a href="#cb66-16" aria-hidden="true" tabindex="-1"></a> <span class="st">"frequent_op"</span>,</span>
<span id="cb66-17"><a href="#cb66-17" aria-hidden="true" tabindex="-1"></a> config<span class="op">=</span>profiling_config</span>
<span id="cb66-18"><a href="#cb66-18" aria-hidden="true" tabindex="-1"></a> ):</span>
<span id="cb66-19"><a href="#cb66-19" aria-hidden="true" tabindex="-1"></a> <span class="co"># This only logs every 10th call, and only if it takes &gt; 1ms</span></span>
<span id="cb66-20"><a href="#cb66-20" aria-hidden="true" tabindex="-1"></a> result <span class="op">=</span> expensive_computation(data)</span>
<span id="cb66-21"><a href="#cb66-21" aria-hidden="true" tabindex="-1"></a> <span class="cf">return</span> result</span></code></pre></div><button title="Copy to Clipboard" class="code-copy-button"><i class="bi"></i></button></div>
<p><strong>ProfilingConfig Parameters</strong>:
- <code>enabled</code>: Enable/disable profiling globally (default: <code>True</code>)
- <code>min_duration_ms</code>: Minimum duration to log in milliseconds (default: <code>0.1</code>)
- <code>log_interval</code>: Log every Nth function call (default: <code>1</code> = log all)</p>
<p><strong>Use cases</strong>:
- <strong>High-frequency methods</strong>: Use <code>log_interval=100</code> to reduce logging overhead
- <strong>Filter noise</strong>: Use <code>min_duration_ms=1.0</code> to skip very fast operations
- <strong>Debugging</strong>: Use <code>log_interval=1, min_duration_ms=0.0</code> to log everything</p>
</section>
<section id="viewing-profiling-metrics" class="level3">
<h3 class="anchored" data-anchor-id="viewing-profiling-metrics">Viewing Profiling Metrics</h3>
<p>In your SwanLab dashboard, profiling metrics appear under the “profiling” namespace:</p>
<pre class="text"><code>profiling/Time taken: AxolotlTrainer.training_step
profiling/Time taken: AxolotlTrainer.prediction_step
profiling/Time taken: MyTrainer.forward_pass
profiling/Time taken: MyTrainer.backward_pass</code></pre>
<p>You can:
- <strong>Track over time</strong>: See if methods get faster/slower during training
- <strong>Compare runs</strong>: Compare profiling metrics across experiments
- <strong>Identify regressions</strong>: Detect if a code change slowed down training</p>
</section>
<section id="configuration-in-axolotl-config" class="level3">
<h3 class="anchored" data-anchor-id="configuration-in-axolotl-config">Configuration in Axolotl Config</h3>
<p>Profiling is automatically enabled when SwanLab is enabled. No additional config needed:</p>
<div class="code-copy-outer-scaffold"><div class="sourceCode" id="cb68"><pre class="sourceCode yaml code-with-copy"><code class="sourceCode yaml"><span id="cb68-1"><a href="#cb68-1" aria-hidden="true" tabindex="-1"></a><span class="fu">plugins</span><span class="kw">:</span></span>
<span id="cb68-2"><a href="#cb68-2" aria-hidden="true" tabindex="-1"></a><span class="at"> </span><span class="kw">-</span><span class="at"> axolotl.integrations.swanlab.SwanLabPlugin</span></span>
<span id="cb68-3"><a href="#cb68-3" aria-hidden="true" tabindex="-1"></a></span>
<span id="cb68-4"><a href="#cb68-4" aria-hidden="true" tabindex="-1"></a><span class="fu">use_swanlab</span><span class="kw">:</span><span class="at"> </span><span class="ch">true</span></span>
<span id="cb68-5"><a href="#cb68-5" aria-hidden="true" tabindex="-1"></a><span class="fu">swanlab_project</span><span class="kw">:</span><span class="at"> my-project</span></span></code></pre></div><button title="Copy to Clipboard" class="code-copy-button"><i class="bi"></i></button></div>
<p>To disable profiling while keeping SwanLab enabled:</p>
<div class="code-copy-outer-scaffold"><div class="sourceCode" id="cb69"><pre class="sourceCode python code-with-copy"><code class="sourceCode python"><span id="cb69-1"><a href="#cb69-1" aria-hidden="true" tabindex="-1"></a><span class="im">from</span> axolotl.integrations.swanlab.profiling <span class="im">import</span> DEFAULT_PROFILING_CONFIG</span>
<span id="cb69-2"><a href="#cb69-2" aria-hidden="true" tabindex="-1"></a></span>
<span id="cb69-3"><a href="#cb69-3" aria-hidden="true" tabindex="-1"></a>DEFAULT_PROFILING_CONFIG.enabled <span class="op">=</span> <span class="va">False</span></span></code></pre></div><button title="Copy to Clipboard" class="code-copy-button"><i class="bi"></i></button></div>
</section>
<section id="performance-impact-1" class="level3">
<h3 class="anchored" data-anchor-id="performance-impact-1">Performance Impact</h3>
<ul>
<li><strong>Decorator overhead</strong>: ~2-5 microseconds per call (negligible)</li>
<li><strong>Context manager overhead</strong>: ~1-3 microseconds (negligible)</li>
<li><strong>Logging overhead</strong>: Only when SwanLab is enabled and method duration exceeds threshold</li>
<li><strong>Network overhead</strong>: SwanLab batches metrics efficiently</li>
</ul>
<p><strong>Expected overhead</strong>: &lt; 0.1% per training step (effectively zero)</p>
</section>
<section id="best-practices" class="level3">
<h3 class="anchored" data-anchor-id="best-practices">Best Practices</h3>
<ol type="1">
<li><strong>Profile bottlenecks first</strong>: Start by profiling suspected slow operations</li>
<li><strong>Use min_duration_ms</strong>: Filter out fast operations (&lt; 1ms) to reduce noise</li>
<li><strong>Throttle high-frequency calls</strong>: Use <code>log_interval</code> for methods called &gt; 100 times/step</li>
<li><strong>Profile across runs</strong>: Compare profiling metrics before/after optimization</li>
<li><strong>Monitor distributed training</strong>: Check for rank-specific slowdowns</li>
</ol>
</section>
<section id="example-complete-profiling-setup" class="level3">
<h3 class="anchored" data-anchor-id="example-complete-profiling-setup">Example: Complete Profiling Setup</h3>
<div class="code-copy-outer-scaffold"><div class="sourceCode" id="cb70"><pre class="sourceCode python code-with-copy"><code class="sourceCode python"><span id="cb70-1"><a href="#cb70-1" aria-hidden="true" tabindex="-1"></a><span class="im">from</span> axolotl.integrations.swanlab.profiling <span class="im">import</span> (</span>
<span id="cb70-2"><a href="#cb70-2" aria-hidden="true" tabindex="-1"></a> swanlab_profile,</span>
<span id="cb70-3"><a href="#cb70-3" aria-hidden="true" tabindex="-1"></a> swanlab_profiling_context,</span>
<span id="cb70-4"><a href="#cb70-4" aria-hidden="true" tabindex="-1"></a> ProfilingConfig,</span>
<span id="cb70-5"><a href="#cb70-5" aria-hidden="true" tabindex="-1"></a>)</span>
<span id="cb70-6"><a href="#cb70-6" aria-hidden="true" tabindex="-1"></a></span>
<span id="cb70-7"><a href="#cb70-7" aria-hidden="true" tabindex="-1"></a><span class="kw">class</span> OptimizedTrainer(AxolotlTrainer):</span>
<span id="cb70-8"><a href="#cb70-8" aria-hidden="true" tabindex="-1"></a> <span class="kw">def</span> <span class="fu">__init__</span>(<span class="va">self</span>, <span class="op">*</span>args, <span class="op">**</span>kwargs):</span>
<span id="cb70-9"><a href="#cb70-9" aria-hidden="true" tabindex="-1"></a> <span class="bu">super</span>().<span class="fu">__init__</span>(<span class="op">*</span>args, <span class="op">**</span>kwargs)</span>
<span id="cb70-10"><a href="#cb70-10" aria-hidden="true" tabindex="-1"></a></span>
<span id="cb70-11"><a href="#cb70-11" aria-hidden="true" tabindex="-1"></a> <span class="co"># Custom profiling config for high-frequency operations</span></span>
<span id="cb70-12"><a href="#cb70-12" aria-hidden="true" tabindex="-1"></a> <span class="va">self</span>.fast_op_config <span class="op">=</span> ProfilingConfig(</span>
<span id="cb70-13"><a href="#cb70-13" aria-hidden="true" tabindex="-1"></a> enabled<span class="op">=</span><span class="va">True</span>,</span>
<span id="cb70-14"><a href="#cb70-14" aria-hidden="true" tabindex="-1"></a> min_duration_ms<span class="op">=</span><span class="fl">0.5</span>,</span>
<span id="cb70-15"><a href="#cb70-15" aria-hidden="true" tabindex="-1"></a> log_interval<span class="op">=</span><span class="dv">50</span>,</span>
<span id="cb70-16"><a href="#cb70-16" aria-hidden="true" tabindex="-1"></a> )</span>
<span id="cb70-17"><a href="#cb70-17" aria-hidden="true" tabindex="-1"></a></span>
<span id="cb70-18"><a href="#cb70-18" aria-hidden="true" tabindex="-1"></a> <span class="at">@swanlab_profile</span></span>
<span id="cb70-19"><a href="#cb70-19" aria-hidden="true" tabindex="-1"></a> <span class="kw">def</span> training_step(<span class="va">self</span>, model, inputs):</span>
<span id="cb70-20"><a href="#cb70-20" aria-hidden="true" tabindex="-1"></a> <span class="co">"""Main training step - always profile."""</span></span>
<span id="cb70-21"><a href="#cb70-21" aria-hidden="true" tabindex="-1"></a> <span class="cf">return</span> <span class="bu">super</span>().training_step(model, inputs)</span>
<span id="cb70-22"><a href="#cb70-22" aria-hidden="true" tabindex="-1"></a></span>
<span id="cb70-23"><a href="#cb70-23" aria-hidden="true" tabindex="-1"></a> <span class="at">@swanlab_profile</span></span>
<span id="cb70-24"><a href="#cb70-24" aria-hidden="true" tabindex="-1"></a> <span class="kw">def</span> compute_loss(<span class="va">self</span>, model, inputs, return_outputs<span class="op">=</span><span class="va">False</span>):</span>
<span id="cb70-25"><a href="#cb70-25" aria-hidden="true" tabindex="-1"></a> <span class="co">"""Loss computation - always profile."""</span></span>
<span id="cb70-26"><a href="#cb70-26" aria-hidden="true" tabindex="-1"></a> <span class="cf">return</span> <span class="bu">super</span>().compute_loss(model, inputs, return_outputs)</span>
<span id="cb70-27"><a href="#cb70-27" aria-hidden="true" tabindex="-1"></a></span>
<span id="cb70-28"><a href="#cb70-28" aria-hidden="true" tabindex="-1"></a> <span class="kw">def</span> _prepare_inputs(<span class="va">self</span>, inputs):</span>
<span id="cb70-29"><a href="#cb70-29" aria-hidden="true" tabindex="-1"></a> <span class="co">"""High-frequency operation - throttled profiling."""</span></span>
<span id="cb70-30"><a href="#cb70-30" aria-hidden="true" tabindex="-1"></a> <span class="cf">with</span> swanlab_profiling_context_advanced(</span>
<span id="cb70-31"><a href="#cb70-31" aria-hidden="true" tabindex="-1"></a> <span class="va">self</span>,</span>
<span id="cb70-32"><a href="#cb70-32" aria-hidden="true" tabindex="-1"></a> <span class="st">"prepare_inputs"</span>,</span>
<span id="cb70-33"><a href="#cb70-33" aria-hidden="true" tabindex="-1"></a> config<span class="op">=</span><span class="va">self</span>.fast_op_config,</span>
<span id="cb70-34"><a href="#cb70-34" aria-hidden="true" tabindex="-1"></a> ):</span>
<span id="cb70-35"><a href="#cb70-35" aria-hidden="true" tabindex="-1"></a> <span class="cf">return</span> <span class="bu">super</span>()._prepare_inputs(inputs)</span></code></pre></div><button title="Copy to Clipboard" class="code-copy-button"><i class="bi"></i></button></div>
</section>
<section id="troubleshooting-1" class="level3">
<h3 class="anchored" data-anchor-id="troubleshooting-1">Troubleshooting</h3>
<section id="profiling-metrics-not-appearing-in-swanlab" class="level4">
<h4 class="anchored" data-anchor-id="profiling-metrics-not-appearing-in-swanlab">Profiling metrics not appearing in SwanLab</h4>
<p><strong>Cause</strong>: SwanLab is not enabled or not initialized.</p>
<p><strong>Solution</strong>:</p>
<div class="code-copy-outer-scaffold"><div class="sourceCode" id="cb71"><pre class="sourceCode yaml code-with-copy"><code class="sourceCode yaml"><span id="cb71-1"><a href="#cb71-1" aria-hidden="true" tabindex="-1"></a><span class="fu">use_swanlab</span><span class="kw">:</span><span class="at"> </span><span class="ch">true</span></span>
<span id="cb71-2"><a href="#cb71-2" aria-hidden="true" tabindex="-1"></a><span class="fu">swanlab_project</span><span class="kw">:</span><span class="at"> my-project</span></span></code></pre></div><button title="Copy to Clipboard" class="code-copy-button"><i class="bi"></i></button></div>
<p>Check logs for:</p>
<pre class="text"><code>INFO: SwanLab initialized for project: my-project</code></pre>
</section>
<section id="too-many-profiling-metrics-cluttering-dashboard" class="level4">
<h4 class="anchored" data-anchor-id="too-many-profiling-metrics-cluttering-dashboard">Too many profiling metrics cluttering dashboard</h4>
<p><strong>Cause</strong>: Profiling every function call for high-frequency operations.</p>
<p><strong>Solution</strong>: Use <code>ProfilingConfig</code> with throttling:</p>
<div class="code-copy-outer-scaffold"><div class="sourceCode" id="cb73"><pre class="sourceCode python code-with-copy"><code class="sourceCode python"><span id="cb73-1"><a href="#cb73-1" aria-hidden="true" tabindex="-1"></a>config <span class="op">=</span> ProfilingConfig(</span>
<span id="cb73-2"><a href="#cb73-2" aria-hidden="true" tabindex="-1"></a> min_duration_ms<span class="op">=</span><span class="fl">1.0</span>, <span class="co"># Skip fast ops</span></span>
<span id="cb73-3"><a href="#cb73-3" aria-hidden="true" tabindex="-1"></a> log_interval<span class="op">=</span><span class="dv">100</span>, <span class="co"># Log every 100th call</span></span>
<span id="cb73-4"><a href="#cb73-4" aria-hidden="true" tabindex="-1"></a>)</span></code></pre></div><button title="Copy to Clipboard" class="code-copy-button"><i class="bi"></i></button></div>
</section>
<section id="profiling-overhead-impacting-training-speed" class="level4">
<h4 class="anchored" data-anchor-id="profiling-overhead-impacting-training-speed">Profiling overhead impacting training speed</h4>
<p><strong>Cause</strong>: Profiling itself should have negligible overhead (&lt; 0.1%). If you see &gt; 1% slowdown, this indicates a bug.</p>
<p><strong>Solution</strong>:
1. Disable profiling temporarily to confirm:
<code>python DEFAULT_PROFILING_CONFIG.enabled = False</code>
2. Report issue with profiling data and trainer details</p>
</section>
<section id="profiling-shows-inconsistent-timing" class="level4">
<h4 class="anchored" data-anchor-id="profiling-shows-inconsistent-timing">Profiling shows inconsistent timing</h4>
<p><strong>Cause</strong>: Normal variation due to GPU warmup, data loading, or system load.</p>
<p><strong>Solution</strong>:
- Ignore first few steps (warmup period)
- Look at average/median timing over many steps
- Use <code>log_interval</code> to reduce noise from individual outliers</p>
</section>
</section>
<section id="complete-config-example" class="level3">
<h3 class="anchored" data-anchor-id="complete-config-example">Complete Config Example</h3>
<p>Here’s a complete example integrating SwanLab with your RVQ-Alpha training:</p>
<div class="code-copy-outer-scaffold"><div class="sourceCode" id="cb74"><pre class="sourceCode yaml code-with-copy"><code class="sourceCode yaml"><span id="cb74-1"><a href="#cb74-1" aria-hidden="true" tabindex="-1"></a><span class="fu">base_model</span><span class="kw">:</span><span class="at"> /path/to/your/model</span></span>
<span id="cb74-2"><a href="#cb74-2" aria-hidden="true" tabindex="-1"></a><span class="fu">model_type</span><span class="kw">:</span><span class="at"> Qwen2ForCausalLM</span></span>
<span id="cb74-3"><a href="#cb74-3" aria-hidden="true" tabindex="-1"></a></span>
<span id="cb74-4"><a href="#cb74-4" aria-hidden="true" tabindex="-1"></a><span class="fu">plugins</span><span class="kw">:</span></span>
<span id="cb74-5"><a href="#cb74-5" aria-hidden="true" tabindex="-1"></a><span class="at"> </span><span class="kw">-</span><span class="at"> axolotl.integrations.swanlab.SwanLabPlugin</span></span>
<span id="cb74-6"><a href="#cb74-6" aria-hidden="true" tabindex="-1"></a><span class="at"> </span><span class="kw">-</span><span class="at"> axolotl.integrations.cut_cross_entropy.CutCrossEntropyPlugin</span></span>
<span id="cb74-7"><a href="#cb74-7" aria-hidden="true" tabindex="-1"></a></span>
<span id="cb74-8"><a href="#cb74-8" aria-hidden="true" tabindex="-1"></a><span class="fu">use_swanlab</span><span class="kw">:</span><span class="at"> </span><span class="ch">true</span></span>
<span id="cb74-9"><a href="#cb74-9" aria-hidden="true" tabindex="-1"></a><span class="fu">swanlab_project</span><span class="kw">:</span><span class="at"> RVQ-Alpha-Training</span></span>
<span id="cb74-10"><a href="#cb74-10" aria-hidden="true" tabindex="-1"></a><span class="fu">swanlab_experiment_name</span><span class="kw">:</span><span class="at"> Qwen2.5-7B-MetaQA-Perturb-P020</span></span>
<span id="cb74-11"><a href="#cb74-11" aria-hidden="true" tabindex="-1"></a><span class="fu">swanlab_description</span><span class="kw">:</span><span class="at"> </span><span class="st">"Training on MetaQA and Perturbation datasets with NEW-RVQ encoding"</span></span>
<span id="cb74-12"><a href="#cb74-12" aria-hidden="true" tabindex="-1"></a><span class="fu">swanlab_mode</span><span class="kw">:</span><span class="at"> cloud</span></span>
<span id="cb74-13"><a href="#cb74-13" aria-hidden="true" tabindex="-1"></a><span class="fu">swanlab_workspace</span><span class="kw">:</span><span class="at"> single-cell-genomics</span></span>
<span id="cb74-14"><a href="#cb74-14" aria-hidden="true" tabindex="-1"></a></span>
<span id="cb74-15"><a href="#cb74-15" aria-hidden="true" tabindex="-1"></a><span class="fu">sequence_len</span><span class="kw">:</span><span class="at"> </span><span class="dv">32768</span></span>
<span id="cb74-16"><a href="#cb74-16" aria-hidden="true" tabindex="-1"></a><span class="fu">micro_batch_size</span><span class="kw">:</span><span class="at"> </span><span class="dv">1</span></span>
<span id="cb74-17"><a href="#cb74-17" aria-hidden="true" tabindex="-1"></a><span class="fu">gradient_accumulation_steps</span><span class="kw">:</span><span class="at"> </span><span class="dv">1</span></span>
<span id="cb74-18"><a href="#cb74-18" aria-hidden="true" tabindex="-1"></a><span class="fu">num_epochs</span><span class="kw">:</span><span class="at"> </span><span class="dv">2</span></span>
<span id="cb74-19"><a href="#cb74-19" aria-hidden="true" tabindex="-1"></a><span class="fu">learning_rate</span><span class="kw">:</span><span class="at"> </span><span class="fl">2e-5</span></span>
<span id="cb74-20"><a href="#cb74-20" aria-hidden="true" tabindex="-1"></a><span class="fu">optimizer</span><span class="kw">:</span><span class="at"> adamw_torch_fused</span></span>
<span id="cb74-21"><a href="#cb74-21" aria-hidden="true" tabindex="-1"></a></span>
<span id="cb74-22"><a href="#cb74-22" aria-hidden="true" tabindex="-1"></a><span class="fu">datasets</span><span class="kw">:</span></span>
<span id="cb74-23"><a href="#cb74-23" aria-hidden="true" tabindex="-1"></a><span class="at"> </span><span class="kw">-</span><span class="at"> </span><span class="fu">path</span><span class="kw">:</span><span class="at"> /path/to/dataset</span></span>
<span id="cb74-24"><a href="#cb74-24" aria-hidden="true" tabindex="-1"></a><span class="at"> </span><span class="fu">type</span><span class="kw">:</span><span class="at"> chat_template</span></span>
<span id="cb74-25"><a href="#cb74-25" aria-hidden="true" tabindex="-1"></a></span>
<span id="cb74-26"><a href="#cb74-26" aria-hidden="true" tabindex="-1"></a><span class="fu">output_dir</span><span class="kw">:</span><span class="at"> ./outputs</span></span></code></pre></div><button title="Copy to Clipboard" class="code-copy-button"><i class="bi"></i></button></div>
</section>
<section id="modes-explained" class="level3">
<h3 class="anchored" data-anchor-id="modes-explained">Modes Explained</h3>
</section>
<section id="cloud-mode-default" class="level3">
<h3 class="anchored" data-anchor-id="cloud-mode-default"><code>cloud</code> Mode (Default)</h3>
<ul>
<li>Syncs experiments to SwanLab cloud in real-time</li>
<li>Requires API key and internet connection</li>
<li>Best for: Team collaboration, remote monitoring</li>
</ul>
</section>
<section id="local-mode" class="level3">
<h3 class="anchored" data-anchor-id="local-mode"><code>local</code> Mode</h3>
<ul>
<li>Saves experiments locally only</li>
<li>No cloud sync</li>
<li>Best for: Local development, air-gapped environments</li>
</ul>
</section>
<section id="offline-mode" class="level3">
<h3 class="anchored" data-anchor-id="offline-mode"><code>offline</code> Mode</h3>
<ul>
<li>Saves metadata locally</li>
<li>Can sync to cloud later using <code>swanlab sync</code></li>
<li>Best for: Unstable internet, sync later</li>
</ul>
</section>
<section id="disabled-mode" class="level3">
<h3 class="anchored" data-anchor-id="disabled-mode"><code>disabled</code> Mode</h3>
<ul>
<li>Turns off SwanLab completely</li>
<li>No logging or tracking</li>
<li>Best for: Debugging, testing</li>
</ul>
</section>
<section id="configuration-validation-conflict-detection" class="level3">
<h3 class="anchored" data-anchor-id="configuration-validation-conflict-detection">Configuration Validation &amp; Conflict Detection</h3>
<p>SwanLab integration includes comprehensive validation and conflict detection to help you catch configuration errors early and avoid performance issues.</p>
</section>
<section id="required-fields-validation" class="level3">
<h3 class="anchored" data-anchor-id="required-fields-validation">Required Fields Validation</h3>
<p>The plugin validates your configuration at startup and provides clear error messages with solutions:</p>
<section id="missing-project-name" class="level4">
<h4 class="anchored" data-anchor-id="missing-project-name">Missing Project Name</h4>
<div class="code-copy-outer-scaffold"><div class="sourceCode" id="cb75"><pre class="sourceCode yaml code-with-copy"><code class="sourceCode yaml"><span id="cb75-1"><a href="#cb75-1" aria-hidden="true" tabindex="-1"></a><span class="fu">use_swanlab</span><span class="kw">:</span><span class="at"> </span><span class="ch">true</span></span></code></pre></div><button title="Copy to Clipboard" class="code-copy-button"><i class="bi"></i></button></div>
<p><strong>Solution</strong>:</p>
<div class="code-copy-outer-scaffold"><div class="sourceCode" id="cb76"><pre class="sourceCode yaml code-with-copy"><code class="sourceCode yaml"><span id="cb76-1"><a href="#cb76-1" aria-hidden="true" tabindex="-1"></a><span class="fu">use_swanlab</span><span class="kw">:</span><span class="at"> </span><span class="ch">true</span></span>
<span id="cb76-2"><a href="#cb76-2" aria-hidden="true" tabindex="-1"></a><span class="fu">swanlab_project</span><span class="kw">:</span><span class="at"> my-project</span></span></code></pre></div><button title="Copy to Clipboard" class="code-copy-button"><i class="bi"></i></button></div>
</section>
<section id="invalid-mode" class="level4">
<h4 class="anchored" data-anchor-id="invalid-mode">Invalid Mode</h4>
<div class="code-copy-outer-scaffold"><div class="sourceCode" id="cb77"><pre class="sourceCode yaml code-with-copy"><code class="sourceCode yaml"><span id="cb77-1"><a href="#cb77-1" aria-hidden="true" tabindex="-1"></a><span class="fu">use_swanlab</span><span class="kw">:</span><span class="at"> </span><span class="ch">true</span></span>
<span id="cb77-2"><a href="#cb77-2" aria-hidden="true" tabindex="-1"></a><span class="fu">swanlab_project</span><span class="kw">:</span><span class="at"> my-project</span></span>
<span id="cb77-3"><a href="#cb77-3" aria-hidden="true" tabindex="-1"></a><span class="fu">swanlab_mode</span><span class="kw">:</span><span class="at"> invalid-mode</span></span></code></pre></div><button title="Copy to Clipboard" class="code-copy-button"><i class="bi"></i></button></div>
<p><strong>Solution</strong>:</p>
<div class="code-copy-outer-scaffold"><div class="sourceCode" id="cb78"><pre class="sourceCode yaml code-with-copy"><code class="sourceCode yaml"><span id="cb78-1"><a href="#cb78-1" aria-hidden="true" tabindex="-1"></a><span class="fu">use_swanlab</span><span class="kw">:</span><span class="at"> </span><span class="ch">true</span></span>
<span id="cb78-2"><a href="#cb78-2" aria-hidden="true" tabindex="-1"></a><span class="fu">swanlab_project</span><span class="kw">:</span><span class="at"> my-project</span></span>
<span id="cb78-3"><a href="#cb78-3" aria-hidden="true" tabindex="-1"></a><span class="fu">swanlab_mode</span><span class="kw">:</span><span class="at"> cloud</span><span class="co"> # or: local, offline, disabled</span></span></code></pre></div><button title="Copy to Clipboard" class="code-copy-button"><i class="bi"></i></button></div>
</section>
<section id="empty-project-name" class="level4">
<h4 class="anchored" data-anchor-id="empty-project-name">Empty Project Name</h4>
<div class="code-copy-outer-scaffold"><div class="sourceCode" id="cb79"><pre class="sourceCode yaml code-with-copy"><code class="sourceCode yaml"><span id="cb79-1"><a href="#cb79-1" aria-hidden="true" tabindex="-1"></a><span class="fu">use_swanlab</span><span class="kw">:</span><span class="at"> </span><span class="ch">true</span></span>
<span id="cb79-2"><a href="#cb79-2" aria-hidden="true" tabindex="-1"></a><span class="fu">swanlab_project</span><span class="kw">:</span><span class="at"> </span><span class="st">""</span></span></code></pre></div><button title="Copy to Clipboard" class="code-copy-button"><i class="bi"></i></button></div>
<p><strong>Solution</strong>:</p>
<div class="code-copy-outer-scaffold"><div class="sourceCode" id="cb80"><pre class="sourceCode yaml code-with-copy"><code class="sourceCode yaml"><span id="cb80-1"><a href="#cb80-1" aria-hidden="true" tabindex="-1"></a><span class="fu">use_swanlab</span><span class="kw">:</span><span class="at"> </span><span class="ch">true</span></span>
<span id="cb80-2"><a href="#cb80-2" aria-hidden="true" tabindex="-1"></a><span class="fu">swanlab_project</span><span class="kw">:</span><span class="at"> my-project</span></span></code></pre></div><button title="Copy to Clipboard" class="code-copy-button"><i class="bi"></i></button></div>
</section>
</section>
<section id="cloud-mode-api-key-warning" class="level3">
<h3 class="anchored" data-anchor-id="cloud-mode-api-key-warning">Cloud Mode API Key Warning</h3>
<p>When using <code>cloud</code> mode without an API key, you’ll receive a warning with multiple solutions:</p>
<div class="code-copy-outer-scaffold"><div class="sourceCode" id="cb81"><pre class="sourceCode yaml code-with-copy"><code class="sourceCode yaml"><span id="cb81-1"><a href="#cb81-1" aria-hidden="true" tabindex="-1"></a><span class="fu">use_swanlab</span><span class="kw">:</span><span class="at"> </span><span class="ch">true</span></span>
<span id="cb81-2"><a href="#cb81-2" aria-hidden="true" tabindex="-1"></a><span class="fu">swanlab_project</span><span class="kw">:</span><span class="at"> my-project</span></span>
<span id="cb81-3"><a href="#cb81-3" aria-hidden="true" tabindex="-1"></a><span class="fu">swanlab_mode</span><span class="kw">:</span><span class="at"> cloud</span></span></code></pre></div><button title="Copy to Clipboard" class="code-copy-button"><i class="bi"></i></button></div>
<p><strong>Solutions</strong>:
1. Set environment variable: <code>export SWANLAB_API_KEY=your-api-key</code>
2. Add to config (less secure): <code>swanlab_api_key: your-api-key</code>
3. Run <code>swanlab login</code> before training
4. Use <code>swanlab_mode: local</code> for offline tracking</p>
</section>
<section id="multi-logger-performance-warnings" class="level3">
<h3 class="anchored" data-anchor-id="multi-logger-performance-warnings">Multi-Logger Performance Warnings</h3>
<p>Using multiple logging tools simultaneously (SwanLab + WandB + MLflow + Comet) can impact training performance:</p>
<section id="two-loggers---warning" class="level4">
<h4 class="anchored" data-anchor-id="two-loggers---warning">Two Loggers - Warning</h4>
<div class="code-copy-outer-scaffold"><div class="sourceCode" id="cb82"><pre class="sourceCode yaml code-with-copy"><code class="sourceCode yaml"><span id="cb82-1"><a href="#cb82-1" aria-hidden="true" tabindex="-1"></a><span class="fu">use_swanlab</span><span class="kw">:</span><span class="at"> </span><span class="ch">true</span></span>
<span id="cb82-2"><a href="#cb82-2" aria-hidden="true" tabindex="-1"></a><span class="fu">swanlab_project</span><span class="kw">:</span><span class="at"> my-project</span></span>
<span id="cb82-3"><a href="#cb82-3" aria-hidden="true" tabindex="-1"></a></span>
<span id="cb82-4"><a href="#cb82-4" aria-hidden="true" tabindex="-1"></a><span class="fu">use_wandb</span><span class="kw">:</span><span class="at"> </span><span class="ch">true</span></span>
<span id="cb82-5"><a href="#cb82-5" aria-hidden="true" tabindex="-1"></a><span class="fu">wandb_project</span><span class="kw">:</span><span class="at"> my-project</span></span></code></pre></div><button title="Copy to Clipboard" class="code-copy-button"><i class="bi"></i></button></div>
<p><strong>Impact</strong>:
- Performance overhead: ~1-2% per logger (cumulative)
- Increased memory usage
- Longer training time per step
- Potential config/callback conflicts</p>
<p><strong>Recommendations</strong>:
- Choose ONE primary logging tool for production training
- Use multiple loggers only for:
- Migration period (transitioning between tools)
- Short comparison runs
- Debugging specific tool issues
- Monitor system resources (CPU, memory) during training</p>
</section>
<section id="three-loggers---error-level-warning" class="level4">
<h4 class="anchored" data-anchor-id="three-loggers---error-level-warning">Three+ Loggers - Error-Level Warning</h4>
<div class="code-copy-outer-scaffold"><div class="sourceCode" id="cb83"><pre class="sourceCode yaml code-with-copy"><code class="sourceCode yaml"><span id="cb83-1"><a href="#cb83-1" aria-hidden="true" tabindex="-1"></a><span class="fu">use_swanlab</span><span class="kw">:</span><span class="at"> </span><span class="ch">true</span></span>
<span id="cb83-2"><a href="#cb83-2" aria-hidden="true" tabindex="-1"></a><span class="fu">swanlab_project</span><span class="kw">:</span><span class="at"> my-project</span></span>
<span id="cb83-3"><a href="#cb83-3" aria-hidden="true" tabindex="-1"></a></span>
<span id="cb83-4"><a href="#cb83-4" aria-hidden="true" tabindex="-1"></a><span class="fu">use_wandb</span><span class="kw">:</span><span class="at"> </span><span class="ch">true</span></span>
<span id="cb83-5"><a href="#cb83-5" aria-hidden="true" tabindex="-1"></a><span class="fu">wandb_project</span><span class="kw">:</span><span class="at"> my-project</span></span>
<span id="cb83-6"><a href="#cb83-6" aria-hidden="true" tabindex="-1"></a></span>
<span id="cb83-7"><a href="#cb83-7" aria-hidden="true" tabindex="-1"></a><span class="fu">use_mlflow</span><span class="kw">:</span><span class="at"> </span><span class="ch">true</span></span>
<span id="cb83-8"><a href="#cb83-8" aria-hidden="true" tabindex="-1"></a><span class="fu">mlflow_tracking_uri</span><span class="kw">:</span><span class="at"> http://localhost:5000</span></span></code></pre></div><button title="Copy to Clipboard" class="code-copy-button"><i class="bi"></i></button></div>
<p><strong>Why This Matters</strong>:
- With 3 loggers: ~4-5% overhead per step → significant slowdown over long training
- Example: 10,000 steps at 2s/step → ~400-500 seconds extra (6-8 minutes)
- Memory overhead scales with number of loggers
- Rare edge cases with callback ordering conflicts</p>
</section>
</section>
<section id="auto-enable-logic" class="level3">
<h3 class="anchored" data-anchor-id="auto-enable-logic">Auto-Enable Logic</h3>
<p>For convenience, SwanLab will auto-enable if you specify a project without setting <code>use_swanlab</code>:</p>
<div class="code-copy-outer-scaffold"><div class="sourceCode" id="cb84"><pre class="sourceCode yaml code-with-copy"><code class="sourceCode yaml"><span id="cb84-1"><a href="#cb84-1" aria-hidden="true" tabindex="-1"></a><span class="fu">swanlab_project</span><span class="kw">:</span><span class="at"> my-project</span></span>
<span id="cb84-2"><a href="#cb84-2" aria-hidden="true" tabindex="-1"></a></span>
<span id="cb84-3"><a href="#cb84-3" aria-hidden="true" tabindex="-1"></a><span class="fu">use_swanlab</span><span class="kw">:</span><span class="at"> </span><span class="ch">true</span></span>
<span id="cb84-4"><a href="#cb84-4" aria-hidden="true" tabindex="-1"></a><span class="fu">swanlab_project</span><span class="kw">:</span><span class="at"> my-project</span></span></code></pre></div><button title="Copy to Clipboard" class="code-copy-button"><i class="bi"></i></button></div>
</section>
<section id="distributed-training-detection" class="level3">
<h3 class="anchored" data-anchor-id="distributed-training-detection">Distributed Training Detection</h3>
<p>In distributed training scenarios (multi-GPU), the plugin automatically detects and reports:</p>
<div class="code-copy-outer-scaffold"><div class="sourceCode" id="cb85"><pre class="sourceCode yaml code-with-copy"><code class="sourceCode yaml"><span id="cb85-1"><a href="#cb85-1" aria-hidden="true" tabindex="-1"></a><span class="fu">use_swanlab</span><span class="kw">:</span><span class="at"> </span><span class="ch">true</span></span>
<span id="cb85-2"><a href="#cb85-2" aria-hidden="true" tabindex="-1"></a><span class="fu">swanlab_project</span><span class="kw">:</span><span class="at"> my-project</span></span>
<span id="cb85-3"><a href="#cb85-3" aria-hidden="true" tabindex="-1"></a><span class="fu">swanlab_mode</span><span class="kw">:</span><span class="at"> cloud</span></span></code></pre></div><button title="Copy to Clipboard" class="code-copy-button"><i class="bi"></i></button></div>
<p><strong>Why Only Rank 0</strong>:
- Avoids duplicate experiment runs
- Reduces network/cloud API overhead on worker ranks
- Prevents race conditions in metric logging</p>
</section>
<section id="authentication" class="level3">
<h3 class="anchored" data-anchor-id="authentication">Authentication</h3>
</section>
<section id="method-1-environment-variable-recommended" class="level3">
<h3 class="anchored" data-anchor-id="method-1-environment-variable-recommended">Method 1: Environment Variable (Recommended)</h3>
<div class="code-copy-outer-scaffold"><div class="sourceCode" id="cb86"><pre class="sourceCode bash code-with-copy"><code class="sourceCode bash"><span id="cb86-1"><a href="#cb86-1" aria-hidden="true" tabindex="-1"></a><span class="bu">export</span> <span class="va">SWANLAB_API_KEY</span><span class="op">=</span>your-api-key-here</span></code></pre></div><button title="Copy to Clipboard" class="code-copy-button"><i class="bi"></i></button></div>
</section>
<section id="method-2-login-command" class="level3">
<h3 class="anchored" data-anchor-id="method-2-login-command">Method 2: Login Command</h3>
<div class="code-copy-outer-scaffold"><div class="sourceCode" id="cb87"><pre class="sourceCode bash code-with-copy"><code class="sourceCode bash"><span id="cb87-1"><a href="#cb87-1" aria-hidden="true" tabindex="-1"></a><span class="ex">swanlab</span> login</span></code></pre></div><button title="Copy to Clipboard" class="code-copy-button"><i class="bi"></i></button></div>
</section>
<section id="method-3-config-file" class="level3">
<h3 class="anchored" data-anchor-id="method-3-config-file">Method 3: Config File</h3>
<div class="code-copy-outer-scaffold"><div class="sourceCode" id="cb88"><pre class="sourceCode yaml code-with-copy"><code class="sourceCode yaml"><span id="cb88-1"><a href="#cb88-1" aria-hidden="true" tabindex="-1"></a><span class="fu">swanlab_api_key</span><span class="kw">:</span><span class="at"> your-api-key-here</span></span></code></pre></div><button title="Copy to Clipboard" class="code-copy-button"><i class="bi"></i></button></div>
</section>
<section id="what-gets-logged" class="level3">
<h3 class="anchored" data-anchor-id="what-gets-logged">What Gets Logged?</h3>
</section>
<section id="automatically-logged-metrics" class="level3">
<h3 class="anchored" data-anchor-id="automatically-logged-metrics">Automatically Logged Metrics</h3>
<ul>
<li>Training loss</li>
<li>Learning rate</li>
<li>Gradient norm</li>
<li>Training steps</li>
<li>Epoch progress</li>
</ul>
</section>
<section id="automatically-logged-config" class="level3">
<h3 class="anchored" data-anchor-id="automatically-logged-config">Automatically Logged Config</h3>
<ul>
<li>Model configuration (base_model, model_type)</li>
<li>Training hyperparameters (learning_rate, batch_size, etc.)</li>
<li>Optimizer settings</li>
<li>Parallelization settings (FSDP, DeepSpeed, Context Parallel)</li>
<li>Axolotl configuration file</li>
<li>DeepSpeed configuration (if used)</li>
</ul>
</section>
<section id="viewing-your-experiments" class="level3">
<h3 class="anchored" data-anchor-id="viewing-your-experiments">Viewing Your Experiments</h3>
</section>
<section id="cloud-mode" class="level3">
<h3 class="anchored" data-anchor-id="cloud-mode">Cloud Mode</h3>
<p>Visit <a href="https://swanlab.cn">https://swanlab.cn</a> and navigate to your project to view:
- Real-time training metrics
- Hyperparameter comparison
- System resource usage
- Configuration files</p>
</section>
<section id="local-mode-1" class="level3">
<h3 class="anchored" data-anchor-id="local-mode-1">Local Mode</h3>
<div class="code-copy-outer-scaffold"><div class="sourceCode" id="cb89"><pre class="sourceCode bash code-with-copy"><code class="sourceCode bash"><span id="cb89-1"><a href="#cb89-1" aria-hidden="true" tabindex="-1"></a><span class="ex">swanlab</span> watch ./swanlog</span></code></pre></div><button title="Copy to Clipboard" class="code-copy-button"><i class="bi"></i></button></div>
</section>
<section id="integration-with-existing-tools" class="level3">
<h3 class="anchored" data-anchor-id="integration-with-existing-tools">Integration with Existing Tools</h3>
<p>SwanLab can work alongside other tracking tools:</p>
<div class="code-copy-outer-scaffold"><div class="sourceCode" id="cb90"><pre class="sourceCode yaml code-with-copy"><code class="sourceCode yaml"><span id="cb90-1"><a href="#cb90-1" aria-hidden="true" tabindex="-1"></a><span class="fu">plugins</span><span class="kw">:</span></span>
<span id="cb90-2"><a href="#cb90-2" aria-hidden="true" tabindex="-1"></a><span class="at"> </span><span class="kw">-</span><span class="at"> axolotl.integrations.swanlab.SwanLabPlugin</span></span>
<span id="cb90-3"><a href="#cb90-3" aria-hidden="true" tabindex="-1"></a></span>
<span id="cb90-4"><a href="#cb90-4" aria-hidden="true" tabindex="-1"></a><span class="fu">use_swanlab</span><span class="kw">:</span><span class="at"> </span><span class="ch">true</span></span>
<span id="cb90-5"><a href="#cb90-5" aria-hidden="true" tabindex="-1"></a><span class="fu">swanlab_project</span><span class="kw">:</span><span class="at"> my-project</span></span>
<span id="cb90-6"><a href="#cb90-6" aria-hidden="true" tabindex="-1"></a></span>
<span id="cb90-7"><a href="#cb90-7" aria-hidden="true" tabindex="-1"></a><span class="fu">use_wandb</span><span class="kw">:</span><span class="at"> </span><span class="ch">true</span></span>
<span id="cb90-8"><a href="#cb90-8" aria-hidden="true" tabindex="-1"></a><span class="fu">wandb_project</span><span class="kw">:</span><span class="at"> my-project</span></span></code></pre></div><button title="Copy to Clipboard" class="code-copy-button"><i class="bi"></i></button></div>
</section>
<section id="troubleshooting-2" class="level3">
<h3 class="anchored" data-anchor-id="troubleshooting-2">Troubleshooting</h3>
</section>
<section id="configuration-errors" class="level3">
<h3 class="anchored" data-anchor-id="configuration-errors">Configuration Errors</h3>
<section id="error-swanlab-enabled-but-swanlab_project-is-not-set" class="level4">
<h4 class="anchored" data-anchor-id="error-swanlab-enabled-but-swanlab_project-is-not-set">Error: “SwanLab enabled but ‘swanlab_project’ is not set”</h4>
<p><strong>Cause</strong>: You enabled SwanLab (<code>use_swanlab: true</code>) but forgot to specify a project name.</p>
<p><strong>Solution</strong>:</p>
<div class="code-copy-outer-scaffold"><div class="sourceCode" id="cb91"><pre class="sourceCode yaml code-with-copy"><code class="sourceCode yaml"><span id="cb91-1"><a href="#cb91-1" aria-hidden="true" tabindex="-1"></a><span class="fu">use_swanlab</span><span class="kw">:</span><span class="at"> </span><span class="ch">true</span></span>
<span id="cb91-2"><a href="#cb91-2" aria-hidden="true" tabindex="-1"></a><span class="fu">swanlab_project</span><span class="kw">:</span><span class="at"> my-project</span><span class="co"> # Add this line</span></span></code></pre></div><button title="Copy to Clipboard" class="code-copy-button"><i class="bi"></i></button></div>
</section>
<section id="error-invalid-swanlab_mode-xxx" class="level4">
<h4 class="anchored" data-anchor-id="error-invalid-swanlab_mode-xxx">Error: “Invalid swanlab_mode: ‘xxx’”</h4>
<p><strong>Cause</strong>: You provided an invalid mode value.</p>
<p><strong>Solution</strong>: Use one of the valid modes:</p>
<div class="code-copy-outer-scaffold"><div class="sourceCode" id="cb92"><pre class="sourceCode yaml code-with-copy"><code class="sourceCode yaml"><span id="cb92-1"><a href="#cb92-1" aria-hidden="true" tabindex="-1"></a><span class="fu">swanlab_mode</span><span class="kw">:</span><span class="at"> cloud</span><span class="co"> # or: local, offline, disabled</span></span></code></pre></div><button title="Copy to Clipboard" class="code-copy-button"><i class="bi"></i></button></div>
</section>
<section id="error-swanlab_project-cannot-be-an-empty-string" class="level4">
<h4 class="anchored" data-anchor-id="error-swanlab_project-cannot-be-an-empty-string">Error: “swanlab_project cannot be an empty string”</h4>
<p><strong>Cause</strong>: You set <code>swanlab_project: ""</code> (empty string).</p>
<p><strong>Solution</strong>: Either provide a valid name or remove the field:</p>
<div class="code-copy-outer-scaffold"><div class="sourceCode" id="cb93"><pre class="sourceCode yaml code-with-copy"><code class="sourceCode yaml"><span id="cb93-1"><a href="#cb93-1" aria-hidden="true" tabindex="-1"></a><span class="fu">swanlab_project</span><span class="kw">:</span><span class="at"> my-project</span></span></code></pre></div><button title="Copy to Clipboard" class="code-copy-button"><i class="bi"></i></button></div>
</section>
</section>
<section id="import-errors" class="level3">
<h3 class="anchored" data-anchor-id="import-errors">Import Errors</h3>
<section id="error-swanlab-is-not-installed" class="level4">
<h4 class="anchored" data-anchor-id="error-swanlab-is-not-installed">Error: “SwanLab is not installed”</h4>
<p><strong>Cause</strong>: SwanLab package is not installed in your environment.</p>
<p><strong>Solution</strong>:</p>
<div class="code-copy-outer-scaffold"><div class="sourceCode" id="cb94"><pre class="sourceCode bash code-with-copy"><code class="sourceCode bash"><span id="cb94-1"><a href="#cb94-1" aria-hidden="true" tabindex="-1"></a><span class="ex">pip</span> install swanlab</span>
<span id="cb94-2"><a href="#cb94-2" aria-hidden="true" tabindex="-1"></a><span class="ex">pip</span> install swanlab<span class="op">&gt;</span>=0.3.0</span></code></pre></div><button title="Copy to Clipboard" class="code-copy-button"><i class="bi"></i></button></div>
</section>
</section>
<section id="performance-issues" class="level3">
<h3 class="anchored" data-anchor-id="performance-issues">Performance Issues</h3>
<section id="warning-multiple-logging-tools-enabled" class="level4">
<h4 class="anchored" data-anchor-id="warning-multiple-logging-tools-enabled">Warning: “Multiple logging tools enabled”</h4>
<p><strong>Cause</strong>: You have multiple experiment tracking tools enabled (e.g., SwanLab + WandB + MLflow).</p>
<p><strong>Impact</strong>: ~1-2% performance overhead per logger, cumulative.</p>
<p><strong>Solution</strong>: For production training, disable all but one logger:</p>
<div class="code-copy-outer-scaffold"><div class="sourceCode" id="cb95"><pre class="sourceCode yaml code-with-copy"><code class="sourceCode yaml"><span id="cb95-1"><a href="#cb95-1" aria-hidden="true" tabindex="-1"></a><span class="fu">use_swanlab</span><span class="kw">:</span><span class="at"> </span><span class="ch">true</span></span>
<span id="cb95-2"><a href="#cb95-2" aria-hidden="true" tabindex="-1"></a><span class="fu">swanlab_project</span><span class="kw">:</span><span class="at"> my-project</span></span>
<span id="cb95-3"><a href="#cb95-3" aria-hidden="true" tabindex="-1"></a><span class="fu">use_wandb</span><span class="kw">:</span><span class="at"> </span><span class="ch">false</span><span class="co"> # Disable others</span></span>
<span id="cb95-4"><a href="#cb95-4" aria-hidden="true" tabindex="-1"></a><span class="fu">use_mlflow</span><span class="kw">:</span><span class="at"> </span><span class="ch">false</span></span>
<span id="cb95-5"><a href="#cb95-5" aria-hidden="true" tabindex="-1"></a></span>
<span id="cb95-6"><a href="#cb95-6" aria-hidden="true" tabindex="-1"></a><span class="fu">use_swanlab</span><span class="kw">:</span><span class="at"> </span><span class="ch">false</span></span>
<span id="cb95-7"><a href="#cb95-7" aria-hidden="true" tabindex="-1"></a><span class="fu">use_wandb</span><span class="kw">:</span><span class="at"> </span><span class="ch">true</span></span>
<span id="cb95-8"><a href="#cb95-8" aria-hidden="true" tabindex="-1"></a><span class="fu">wandb_project</span><span class="kw">:</span><span class="at"> my-project</span></span></code></pre></div><button title="Copy to Clipboard" class="code-copy-button"><i class="bi"></i></button></div>
<p><strong>Exception</strong>: Multiple loggers are acceptable for:
- Short comparison runs (&lt; 100 steps)
- Migration testing between logging tools
- Debugging logger-specific issues</p>
</section>
</section>
<section id="distributed-training-issues" class="level3">
<h3 class="anchored" data-anchor-id="distributed-training-issues">Distributed Training Issues</h3>
<section id="swanlab-creates-duplicate-runs-in-multi-gpu-training" class="level4">
<h4 class="anchored" data-anchor-id="swanlab-creates-duplicate-runs-in-multi-gpu-training">SwanLab creates duplicate runs in multi-GPU training</h4>
<p><strong>Cause</strong>: All ranks are initializing SwanLab instead of just rank 0.</p>
<p><strong>Expected Behavior</strong>: The plugin automatically ensures only rank 0 initializes SwanLab. You should see:</p>
<pre class="text"><code>Info: Distributed training detected (world_size=4)
Info: Only rank 0 will initialize SwanLab
Info: Other ranks will skip SwanLab to avoid conflicts</code></pre>
<p><strong>If you see duplicates</strong>:
1. Check your plugin is loaded correctly
2. Verify you’re using the latest SwanLab integration code
3. Check logs for initialization messages on all ranks</p>
</section>
</section>
<section id="swanlab-not-logging-metrics" class="level3">
<h3 class="anchored" data-anchor-id="swanlab-not-logging-metrics">SwanLab not logging metrics</h3>
<p><strong>Solution</strong>: Ensure SwanLab is initialized before training starts. The plugin automatically handles this in <code>pre_model_load</code>.</p>
</section>
<section id="api-key-errors" class="level3">
<h3 class="anchored" data-anchor-id="api-key-errors">API Key errors</h3>
<p><strong>Solution</strong>:</p>
<div class="code-copy-outer-scaffold"><div class="sourceCode" id="cb97"><pre class="sourceCode bash code-with-copy"><code class="sourceCode bash"><span id="cb97-1"><a href="#cb97-1" aria-hidden="true" tabindex="-1"></a><span class="bu">echo</span> <span class="va">$SWANLAB_API_KEY</span></span>
<span id="cb97-2"><a href="#cb97-2" aria-hidden="true" tabindex="-1"></a></span>
<span id="cb97-3"><a href="#cb97-3" aria-hidden="true" tabindex="-1"></a><span class="ex">swanlab</span> login</span></code></pre></div><button title="Copy to Clipboard" class="code-copy-button"><i class="bi"></i></button></div>
</section>
<section id="cloud-sync-issues" class="level3">
<h3 class="anchored" data-anchor-id="cloud-sync-issues">Cloud sync issues</h3>
<p><strong>Solution</strong>: Use <code>offline</code> mode and sync later:</p>
<div class="code-copy-outer-scaffold"><div class="sourceCode" id="cb98"><pre class="sourceCode yaml code-with-copy"><code class="sourceCode yaml"><span id="cb98-1"><a href="#cb98-1" aria-hidden="true" tabindex="-1"></a><span class="fu">swanlab_mode</span><span class="kw">:</span><span class="at"> offline</span></span></code></pre></div><button title="Copy to Clipboard" class="code-copy-button"><i class="bi"></i></button></div>
<p>Then sync when ready:</p>
<div class="code-copy-outer-scaffold"><div class="sourceCode" id="cb99"><pre class="sourceCode bash code-with-copy"><code class="sourceCode bash"><span id="cb99-1"><a href="#cb99-1" aria-hidden="true" tabindex="-1"></a><span class="ex">swanlab</span> sync ./swanlog</span></code></pre></div><button title="Copy to Clipboard" class="code-copy-button"><i class="bi"></i></button></div>
</section>
<section id="plugin-not-loaded" class="level3">
<h3 class="anchored" data-anchor-id="plugin-not-loaded">Plugin not loaded</h3>
<p><strong>Solution</strong>: Verify plugin path in config:</p>
<div class="code-copy-outer-scaffold"><div class="sourceCode" id="cb100"><pre class="sourceCode yaml code-with-copy"><code class="sourceCode yaml"><span id="cb100-1"><a href="#cb100-1" aria-hidden="true" tabindex="-1"></a><span class="fu">plugins</span><span class="kw">:</span></span>
<span id="cb100-2"><a href="#cb100-2" aria-hidden="true" tabindex="-1"></a><span class="at"> </span><span class="kw">-</span><span class="at"> axolotl.integrations.swanlab.SwanLabPlugin</span><span class="co"> # Correct path</span></span></code></pre></div><button title="Copy to Clipboard" class="code-copy-button"><i class="bi"></i></button></div>
</section>
<section id="lark-notification-issues" class="level3">
<h3 class="anchored" data-anchor-id="lark-notification-issues">Lark Notification Issues</h3>
<section id="error-failed-to-import-swanlab-lark-plugin" class="level4">
<h4 class="anchored" data-anchor-id="error-failed-to-import-swanlab-lark-plugin">Error: “Failed to import SwanLab Lark plugin”</h4>
<p><strong>Cause</strong>: Your SwanLab version doesn’t include the Lark plugin (requires SwanLab &gt;= 0.3.0).</p>
<p><strong>Solution</strong>:</p>
<div class="code-copy-outer-scaffold"><div class="sourceCode" id="cb101"><pre class="sourceCode bash code-with-copy"><code class="sourceCode bash"><span id="cb101-1"><a href="#cb101-1" aria-hidden="true" tabindex="-1"></a><span class="ex">pip</span> install <span class="at">--upgrade</span> swanlab</span>
<span id="cb101-2"><a href="#cb101-2" aria-hidden="true" tabindex="-1"></a></span>
<span id="cb101-3"><a href="#cb101-3" aria-hidden="true" tabindex="-1"></a><span class="ex">pip</span> install <span class="st">'swanlab&gt;=0.3.0'</span></span></code></pre></div><button title="Copy to Clipboard" class="code-copy-button"><i class="bi"></i></button></div>
</section>
<section id="warning-lark-webhook-has-no-secret-configured" class="level4">
<h4 class="anchored" data-anchor-id="warning-lark-webhook-has-no-secret-configured">Warning: “Lark webhook has no secret configured”</h4>
<p><strong>Cause</strong>: You provided <code>swanlab_lark_webhook_url</code> but no <code>swanlab_lark_secret</code>.</p>
<p><strong>Impact</strong>: Lark notifications will work, but without HMAC authentication (security risk).</p>
<p><strong>Solution</strong>: Add HMAC secret for production use:</p>
<div class="code-copy-outer-scaffold"><div class="sourceCode" id="cb102"><pre class="sourceCode yaml code-with-copy"><code class="sourceCode yaml"><span id="cb102-1"><a href="#cb102-1" aria-hidden="true" tabindex="-1"></a><span class="fu">swanlab_lark_webhook_url</span><span class="kw">:</span><span class="at"> https://open.feishu.cn/open-apis/bot/v2/hook/xxx</span></span>
<span id="cb102-2"><a href="#cb102-2" aria-hidden="true" tabindex="-1"></a><span class="fu">swanlab_lark_secret</span><span class="kw">:</span><span class="at"> your-webhook-secret</span><span class="co"> # Add this line</span></span></code></pre></div><button title="Copy to Clipboard" class="code-copy-button"><i class="bi"></i></button></div>
<p><strong>When it’s OK to skip secret</strong>:
- Local development and testing
- Internal networks with restricted access
- Non-sensitive training experiments</p>
<p><strong>When secret is required</strong>:
- Production training jobs
- Training with proprietary data
- Multi-team shared Lark groups</p>
</section>
<section id="error-failed-to-register-lark-callback" class="level4">
<h4 class="anchored" data-anchor-id="error-failed-to-register-lark-callback">Error: “Failed to register Lark callback”</h4>
<p><strong>Cause</strong>: Invalid webhook URL or network connectivity issues.</p>
<p><strong>Diagnostic steps</strong>:</p>
<div class="code-copy-outer-scaffold"><div class="sourceCode" id="cb103"><pre class="sourceCode bash code-with-copy"><code class="sourceCode bash"><span id="cb103-1"><a href="#cb103-1" aria-hidden="true" tabindex="-1"></a><span class="ex">curl</span> <span class="at">-X</span> POST <span class="st">"YOUR_WEBHOOK_URL"</span> <span class="dt">\</span></span>
<span id="cb103-2"><a href="#cb103-2" aria-hidden="true" tabindex="-1"></a> <span class="at">-H</span> <span class="st">'Content-Type: application/json'</span> <span class="dt">\</span></span>
<span id="cb103-3"><a href="#cb103-3" aria-hidden="true" tabindex="-1"></a> <span class="at">-d</span> <span class="st">'{"msg_type":"text","content":{"text":"Test from Axolotl"}}'</span></span>
<span id="cb103-4"><a href="#cb103-4" aria-hidden="true" tabindex="-1"></a></span>
<span id="cb103-5"><a href="#cb103-5" aria-hidden="true" tabindex="-1"></a><span class="ex">pip</span> show swanlab</span></code></pre></div><button title="Copy to Clipboard" class="code-copy-button"><i class="bi"></i></button></div>
<p><strong>Solution</strong>:
1. Verify webhook URL is correct (copy from Lark bot settings)
2. Check network connectivity to Lark API
3. Ensure webhook is not expired (Lark webhooks can expire)
4. Regenerate webhook URL in Lark bot settings if needed</p>
</section>
<section id="lark-notifications-not-received" class="level4">
<h4 class="anchored" data-anchor-id="lark-notifications-not-received">Lark notifications not received</h4>
<p><strong>Cause</strong>: Multiple possible causes.</p>
<p><strong>Diagnostic checklist</strong>:</p>
<ol type="1">
<li><p><strong>Check training logs</strong> for Lark registration confirmation:</p>
<pre class="text"><code># Expected log message (rank 0 only):
INFO: Registered Lark notification callback with HMAC authentication</code></pre></li>
<li><p><strong>Verify webhook in Lark</strong>: Test webhook manually (see above)</p></li>
<li><p><strong>Check distributed training</strong>: Only rank 0 sends notifications</p>
<div class="code-copy-outer-scaffold"><div class="sourceCode" id="cb105"><pre class="sourceCode bash code-with-copy"><code class="sourceCode bash"><span id="cb105-1"><a href="#cb105-1" aria-hidden="true" tabindex="-1"></a><span class="co"># If running multi-GPU, check rank 0 logs specifically</span></span>
<span id="cb105-2"><a href="#cb105-2" aria-hidden="true" tabindex="-1"></a><span class="fu">grep</span> <span class="st">"Registered Lark"</span> logs/rank_0.log</span></code></pre></div><button title="Copy to Clipboard" class="code-copy-button"><i class="bi"></i></button></div></li>
<li><p><strong>Verify SwanLab is initialized</strong>: Lark callback needs SwanLab to be running</p>
<div class="code-copy-outer-scaffold"><div class="sourceCode" id="cb106"><pre class="sourceCode yaml code-with-copy"><code class="sourceCode yaml"><span id="cb106-1"><a href="#cb106-1" aria-hidden="true" tabindex="-1"></a><span class="fu">use_swanlab</span><span class="kw">:</span><span class="at"> </span><span class="ch">true</span><span class="co"> # Must be enabled</span></span>
<span id="cb106-2"><a href="#cb106-2" aria-hidden="true" tabindex="-1"></a><span class="fu">swanlab_project</span><span class="kw">:</span><span class="at"> my-project</span><span class="co"> # Must be set</span></span></code></pre></div><button title="Copy to Clipboard" class="code-copy-button"><i class="bi"></i></button></div></li>
<li><p><strong>Check Lark bot permissions</strong>: Ensure bot is added to the target group chat</p></li>
</ol>
</section>
<section id="duplicate-lark-notifications-in-multi-gpu-training" class="level4">
<h4 class="anchored" data-anchor-id="duplicate-lark-notifications-in-multi-gpu-training">Duplicate Lark notifications in multi-GPU training</h4>
<p><strong>Expected Behavior</strong>: Should NOT happen - only rank 0 sends notifications.</p>
<p><strong>If you see duplicates</strong>:
1. Check that all GPUs are using the same config file
2. Verify plugin is loaded correctly on all ranks
3. Check logs for unexpected Lark initialization on non-zero ranks
4. Ensure <code>RANK</code> or <code>LOCAL_RANK</code> environment variables are set correctly</p>
<p><strong>Solution</strong>: This is a bug if it occurs. Report with:
- Full training command
- Logs from all ranks
- Config file</p>
</section>
</section>
<section id="comparison-swanlab-vs-wandb" class="level3">
<h3 class="anchored" data-anchor-id="comparison-swanlab-vs-wandb">Comparison: SwanLab vs WandB</h3>
<table class="caption-top table">
<thead>
<tr class="header">
<th>Feature</th>
<th>SwanLab</th>
<th>WandB</th>
</tr>
</thead>
<tbody>
<tr class="odd">
<td>Open Source</td>
<td>✅ Yes</td>
<td>❌ No</td>
</tr>
<tr class="even">
<td>Self-Hosting</td>
<td>✅ Easy</td>
<td>⚠️ Complex</td>
</tr>
<tr class="odd">
<td>Free Tier</td>
<td>✅ Generous</td>
<td>⚠️ Limited</td>
</tr>
<tr class="even">
<td>Chinese Support</td>
<td>✅ Native</td>
<td>⚠️ Limited</td>
</tr>
<tr class="odd">
<td>Offline Mode</td>
<td>✅ Full support</td>
<td>✅ Supported</td>
</tr>
<tr class="even">
<td>Integration</td>
<td>🆕 New</td>
<td>✅ Mature</td>
</tr>
</tbody>
</table>
</section>
<section id="advanced-usage" class="level3">
<h3 class="anchored" data-anchor-id="advanced-usage">Advanced Usage</h3>
</section>
<section id="custom-logging" class="level3">
<h3 class="anchored" data-anchor-id="custom-logging">Custom Logging</h3>
<p>You can add custom metrics in your callbacks:</p>
<div class="code-copy-outer-scaffold"><div class="sourceCode" id="cb107"><pre class="sourceCode python code-with-copy"><code class="sourceCode python"><span id="cb107-1"><a href="#cb107-1" aria-hidden="true" tabindex="-1"></a><span class="im">import</span> swanlab</span>
<span id="cb107-2"><a href="#cb107-2" aria-hidden="true" tabindex="-1"></a></span>
<span id="cb107-3"><a href="#cb107-3" aria-hidden="true" tabindex="-1"></a>swanlab.log({</span>
<span id="cb107-4"><a href="#cb107-4" aria-hidden="true" tabindex="-1"></a> <span class="st">"custom_metric"</span>: value,</span>
<span id="cb107-5"><a href="#cb107-5" aria-hidden="true" tabindex="-1"></a> <span class="st">"epoch"</span>: epoch_num</span>
<span id="cb107-6"><a href="#cb107-6" aria-hidden="true" tabindex="-1"></a>})</span></code></pre></div><button title="Copy to Clipboard" class="code-copy-button"><i class="bi"></i></button></div>
</section>
<section id="experiment-comparison" class="level3">
<h3 class="anchored" data-anchor-id="experiment-comparison">Experiment Comparison</h3>
<div class="code-copy-outer-scaffold"><div class="sourceCode" id="cb108"><pre class="sourceCode bash code-with-copy"><code class="sourceCode bash"><span id="cb108-1"><a href="#cb108-1" aria-hidden="true" tabindex="-1"></a><span class="ex">swanlab</span> compare run1 run2 run3</span></code></pre></div><button title="Copy to Clipboard" class="code-copy-button"><i class="bi"></i></button></div>
</section>
<section id="support" class="level3">
<h3 class="anchored" data-anchor-id="support">Support</h3>
<ul>
<li><strong>Documentation</strong>: <a href="https://docs.swanlab.cn">https://docs.swanlab.cn</a></li>
<li><strong>GitHub</strong>: <a href="https://github.com/SwanHubX/SwanLab">https://github.com/SwanHubX/SwanLab</a></li>
<li><strong>Issues</strong>: Report bugs at <a href="https://github.com/SwanHubX/SwanLab/issues">GitHub Issues</a></li>
</ul>
</section>
<section id="license" class="level3">
<h3 class="anchored" data-anchor-id="license">License</h3>
<p>This integration follows the Axolotl Community License Agreement.</p>
</section>
<section id="acknowledgements" class="level3">
<h3 class="anchored" data-anchor-id="acknowledgements">Acknowledgements</h3>
<p>This integration is built on top of:
- <a href="https://github.com/SwanHubX/SwanLab">SwanLab</a> - Experiment tracking tool
- <a href="https://github.com/huggingface/transformers">Transformers</a> - SwanLabCallback
- <a href="https://github.com/axolotl-ai-cloud/axolotl">Axolotl</a> - Training framework</p>
<p>Please see reference <a href="https://github.com/axolotl-ai-cloud/axolotl/tree/main/src/axolotl/integrations/swanlab">here</a></p>
</section>
</section>
<section id="adding-a-new-integration" class="level2">
<h2 class="anchored" data-anchor-id="adding-a-new-integration">Adding a new integration</h2>
<p>Plugins can be used to customize the behavior of the training pipeline through <a href="https://en.wikipedia.org/wiki/Hooking">hooks</a>. See <a href="https://github.com/axolotl-ai-cloud/axolotl/blob/main/src/axolotl/integrations/base.py"><code>axolotl.integrations.BasePlugin</code></a> for the possible hooks.</p>
<p>To add a new integration, please follow these steps:</p>
<ol type="1">
<li>Create a new folder in the <code>src/axolotl/integrations</code> directory.</li>
<li>Add any relevant files (<code>LICENSE</code>, <code>README.md</code>, <code>ACKNOWLEDGEMENTS.md</code>, etc.) to the new folder.</li>
<li>Add <code>__init__.py</code> and <code>args.py</code> files to the new folder.</li>
</ol>
<ul>
<li><code>__init__.py</code> should import the integration and hook into the appropriate functions.</li>
<li><code>args.py</code> should define the arguments for the integration.</li>
</ul>
<ol start="4" type="1">
<li>(If applicable) Add CPU tests under <code>tests/integrations</code> or GPU tests under <code>tests/e2e/integrations</code>.</li>
</ol>
<div class="callout callout-style-default callout-tip callout-titled">
<div class="callout-header d-flex align-content-center">
<div class="callout-icon-container">
<i class="callout-icon"></i>
</div>
<div class="callout-title-container flex-fill">
Tip
</div>
</div>
<div class="callout-body-container callout-body">
<p>See <a href="https://github.com/axolotl-ai-cloud/axolotl/tree/main/src/axolotl/integrations/cut_cross_entropy">src/axolotl/integrations/cut_cross_entropy</a> for a minimal integration example.</p>
</div>
</div>
<div class="callout callout-style-default callout-warning callout-titled">
<div class="callout-header d-flex align-content-center">
<div class="callout-icon-container">
<i class="callout-icon"></i>
</div>
<div class="callout-title-container flex-fill">
Warning
</div>
</div>
<div class="callout-body-container callout-body">
<p>If you could not load your integration, please ensure you are pip installing in editable mode.</p>
<div class="code-copy-outer-scaffold"><div class="sourceCode" id="cb109"><pre class="sourceCode bash code-with-copy"><code class="sourceCode bash"><span id="cb109-1"><a href="#cb109-1" aria-hidden="true" tabindex="-1"></a><span class="ex">pip</span> install <span class="at">-e</span> .</span></code></pre></div><button title="Copy to Clipboard" class="code-copy-button"><i class="bi"></i></button></div>
<p>and correctly spelled the integration name in the config file.</p>
<div class="code-copy-outer-scaffold"><div class="sourceCode" id="cb110"><pre class="sourceCode yaml code-with-copy"><code class="sourceCode yaml"><span id="cb110-1"><a href="#cb110-1" aria-hidden="true" tabindex="-1"></a><span class="fu">plugins</span><span class="kw">:</span></span>
<span id="cb110-2"><a href="#cb110-2" aria-hidden="true" tabindex="-1"></a><span class="at"> </span><span class="kw">-</span><span class="at"> axolotl.integrations.your_integration_name.YourIntegrationPlugin</span></span></code></pre></div><button title="Copy to Clipboard" class="code-copy-button"><i class="bi"></i></button></div>
</div>
</div>
<div class="callout callout-style-default callout-note callout-titled">
<div class="callout-header d-flex align-content-center">
<div class="callout-icon-container">
<i class="callout-icon"></i>
</div>
<div class="callout-title-container flex-fill">
Note
</div>
</div>
<div class="callout-body-container callout-body">
<p>It is not necessary to place your integration in the <code>integrations</code> folder. It can be in any location, so long as it’s installed in a package in your python env.</p>
<p>See this repo for an example: <a href="https://github.com/axolotl-ai-cloud/diff-transformer">https://github.com/axolotl-ai-cloud/diff-transformer</a></p>
</div>
</div>
</section>
</main> <!-- /main -->
<script id="quarto-html-after-body" type="application/javascript">
window.document.addEventListener("DOMContentLoaded", function (event) {
const icon = "";
const anchorJS = new window.AnchorJS();
anchorJS.options = {
placement: 'right',
icon: icon
};
anchorJS.add('.anchored');
const isCodeAnnotation = (el) => {
for (const clz of el.classList) {
if (clz.startsWith('code-annotation-')) {
return true;
}
}
return false;
}
const onCopySuccess = function(e) {
// button target
const button = e.trigger;
// don't keep focus
button.blur();
// flash "checked"
button.classList.add('code-copy-button-checked');
var currentTitle = button.getAttribute("title");
button.setAttribute("title", "Copied!");
let tooltip;
if (window.bootstrap) {
button.setAttribute("data-bs-toggle", "tooltip");
button.setAttribute("data-bs-placement", "left");
button.setAttribute("data-bs-title", "Copied!");
tooltip = new bootstrap.Tooltip(button,
{ trigger: "manual",
customClass: "code-copy-button-tooltip",
offset: [0, -8]});
tooltip.show();
}
setTimeout(function() {
if (tooltip) {
tooltip.hide();
button.removeAttribute("data-bs-title");
button.removeAttribute("data-bs-toggle");
button.removeAttribute("data-bs-placement");
}
button.setAttribute("title", currentTitle);
button.classList.remove('code-copy-button-checked');
}, 1000);
// clear code selection
e.clearSelection();
}
const getTextToCopy = function(trigger) {
const outerScaffold = trigger.parentElement.cloneNode(true);
const codeEl = outerScaffold.querySelector('code');
for (const childEl of codeEl.children) {
if (isCodeAnnotation(childEl)) {
childEl.remove();
}
}
return codeEl.innerText;
}
const clipboard = new window.ClipboardJS('.code-copy-button:not([data-in-quarto-modal])', {
text: getTextToCopy
});
clipboard.on('success', onCopySuccess);
if (window.document.getElementById('quarto-embedded-source-code-modal')) {
const clipboardModal = new window.ClipboardJS('.code-copy-button[data-in-quarto-modal]', {
text: getTextToCopy,
container: window.document.getElementById('quarto-embedded-source-code-modal')
});
clipboardModal.on('success', onCopySuccess);
}
var localhostRegex = new RegExp(/^(?:http|https):\/\/localhost\:?[0-9]*\//);
var mailtoRegex = new RegExp(/^mailto:/);
var filterRegex = new RegExp("https:\/\/docs\.axolotl\.ai");
var isInternal = (href) => {
return filterRegex.test(href) || localhostRegex.test(href) || mailtoRegex.test(href);
}
// Inspect non-navigation links and adorn them if external
var links = window.document.querySelectorAll('a[href]:not(.nav-link):not(.navbar-brand):not(.toc-action):not(.sidebar-link):not(.sidebar-item-toggle):not(.pagination-link):not(.no-external):not([aria-hidden]):not(.dropdown-item):not(.quarto-navigation-tool):not(.about-link)');
for (var i=0; i<links.length; i++) {
const link = links[i];
if (!isInternal(link.href)) {
// undo the damage that might have been done by quarto-nav.js in the case of
// links that we want to consider external
if (link.dataset.originalHref !== undefined) {
link.href = link.dataset.originalHref;
}
}
}
function tippyHover(el, contentFn, onTriggerFn, onUntriggerFn) {
const config = {
allowHTML: true,
maxWidth: 500,
delay: 100,
arrow: false,
appendTo: function(el) {
return el.parentElement;
},
interactive: true,
interactiveBorder: 10,
theme: 'quarto',
placement: 'bottom-start',
};
if (contentFn) {
config.content = contentFn;
}
if (onTriggerFn) {
config.onTrigger = onTriggerFn;
}
if (onUntriggerFn) {
config.onUntrigger = onUntriggerFn;
}
window.tippy(el, config);
}
const noterefs = window.document.querySelectorAll('a[role="doc-noteref"]');
for (var i=0; i<noterefs.length; i++) {
const ref = noterefs[i];
tippyHover(ref, function() {
// use id or data attribute instead here
let href = ref.getAttribute('data-footnote-href') || ref.getAttribute('href');
try { href = new URL(href).hash; } catch {}
const id = href.replace(/^#\/?/, "");
const note = window.document.getElementById(id);
if (note) {
return note.innerHTML;
} else {
return "";
}
});
}
const xrefs = window.document.querySelectorAll('a.quarto-xref');
const processXRef = (id, note) => {
// Strip column container classes
const stripColumnClz = (el) => {
el.classList.remove("page-full", "page-columns");
if (el.children) {
for (const child of el.children) {
stripColumnClz(child);
}
}
}
stripColumnClz(note)
if (id === null || id.startsWith('sec-')) {
// Special case sections, only their first couple elements
const container = document.createElement("div");
if (note.children && note.children.length > 2) {
container.appendChild(note.children[0].cloneNode(true));
for (let i = 1; i < note.children.length; i++) {
const child = note.children[i];
if (child.tagName === "P" && child.innerText === "") {
continue;
} else {
container.appendChild(child.cloneNode(true));
break;
}
}
if (window.Quarto?.typesetMath) {
window.Quarto.typesetMath(container);
}
return container.innerHTML
} else {
if (window.Quarto?.typesetMath) {
window.Quarto.typesetMath(note);
}
return note.innerHTML;
}
} else {
// Remove any anchor links if they are present
const anchorLink = note.querySelector('a.anchorjs-link');
if (anchorLink) {
anchorLink.remove();
}
if (window.Quarto?.typesetMath) {
window.Quarto.typesetMath(note);
}
if (note.classList.contains("callout")) {
return note.outerHTML;
} else {
return note.innerHTML;
}
}
}
for (var i=0; i<xrefs.length; i++) {
const xref = xrefs[i];
tippyHover(xref, undefined, function(instance) {
instance.disable();
let url = xref.getAttribute('href');
let hash = undefined;
if (url.startsWith('#')) {
hash = url;
} else {
try { hash = new URL(url).hash; } catch {}
}
if (hash) {
const id = hash.replace(/^#\/?/, "");
const note = window.document.getElementById(id);
if (note !== null) {
try {
const html = processXRef(id, note.cloneNode(true));
instance.setContent(html);
} finally {
instance.enable();
instance.show();
}
} else {
// See if we can fetch this
fetch(url.split('#')[0])
.then(res => res.text())
.then(html => {
const parser = new DOMParser();
const htmlDoc = parser.parseFromString(html, "text/html");
const note = htmlDoc.getElementById(id);
if (note !== null) {
const html = processXRef(id, note);
instance.setContent(html);
}
}).finally(() => {
instance.enable();
instance.show();
});
}
} else {
// See if we can fetch a full url (with no hash to target)
// This is a special case and we should probably do some content thinning / targeting
fetch(url)
.then(res => res.text())
.then(html => {
const parser = new DOMParser();
const htmlDoc = parser.parseFromString(html, "text/html");
const note = htmlDoc.querySelector('main.content');
if (note !== null) {
// This should only happen for chapter cross references
// (since there is no id in the URL)
// remove the first header
if (note.children.length > 0 && note.children[0].tagName === "HEADER") {
note.children[0].remove();
}
const html = processXRef(null, note);
instance.setContent(html);
}
}).finally(() => {
instance.enable();
instance.show();
});
}
}, function(instance) {
});
}
let selectedAnnoteEl;
const selectorForAnnotation = ( cell, annotation) => {
let cellAttr = 'data-code-cell="' + cell + '"';
let lineAttr = 'data-code-annotation="' + annotation + '"';
const selector = 'span[' + cellAttr + '][' + lineAttr + ']';
return selector;
}
const selectCodeLines = (annoteEl) => {
const doc = window.document;
const targetCell = annoteEl.getAttribute("data-target-cell");
const targetAnnotation = annoteEl.getAttribute("data-target-annotation");
const annoteSpan = window.document.querySelector(selectorForAnnotation(targetCell, targetAnnotation));
const lines = annoteSpan.getAttribute("data-code-lines").split(",");
const lineIds = lines.map((line) => {
return targetCell + "-" + line;
})
let top = null;
let height = null;
let parent = null;
if (lineIds.length > 0) {
//compute the position of the single el (top and bottom and make a div)
const el = window.document.getElementById(lineIds[0]);
top = el.offsetTop;
height = el.offsetHeight;
parent = el.parentElement.parentElement;
if (lineIds.length > 1) {
const lastEl = window.document.getElementById(lineIds[lineIds.length - 1]);
const bottom = lastEl.offsetTop + lastEl.offsetHeight;
height = bottom - top;
}
if (top !== null && height !== null && parent !== null) {
// cook up a div (if necessary) and position it
let div = window.document.getElementById("code-annotation-line-highlight");
if (div === null) {
div = window.document.createElement("div");
div.setAttribute("id", "code-annotation-line-highlight");
div.style.position = 'absolute';
parent.appendChild(div);
}
div.style.top = top - 2 + "px";
div.style.height = height + 4 + "px";
div.style.left = 0;
let gutterDiv = window.document.getElementById("code-annotation-line-highlight-gutter");
if (gutterDiv === null) {
gutterDiv = window.document.createElement("div");
gutterDiv.setAttribute("id", "code-annotation-line-highlight-gutter");
gutterDiv.style.position = 'absolute';
const codeCell = window.document.getElementById(targetCell);
const gutter = codeCell.querySelector('.code-annotation-gutter');
gutter.appendChild(gutterDiv);
}
gutterDiv.style.top = top - 2 + "px";
gutterDiv.style.height = height + 4 + "px";
}
selectedAnnoteEl = annoteEl;
}
};
const unselectCodeLines = () => {
const elementsIds = ["code-annotation-line-highlight", "code-annotation-line-highlight-gutter"];
elementsIds.forEach((elId) => {
const div = window.document.getElementById(elId);
if (div) {
div.remove();
}
});
selectedAnnoteEl = undefined;
};
// Handle positioning of the toggle
window.addEventListener(
"resize",
throttle(() => {
elRect = undefined;
if (selectedAnnoteEl) {
selectCodeLines(selectedAnnoteEl);
}
}, 10)
);
function throttle(fn, ms) {
let throttle = false;
let timer;
return (...args) => {
if(!throttle) { // first call gets through
fn.apply(this, args);
throttle = true;
} else { // all the others get throttled
if(timer) clearTimeout(timer); // cancel #2
timer = setTimeout(() => {
fn.apply(this, args);
timer = throttle = false;
}, ms);
}
};
}
// Attach click handler to the DT
const annoteDls = window.document.querySelectorAll('dt[data-target-cell]');
for (const annoteDlNode of annoteDls) {
annoteDlNode.addEventListener('click', (event) => {
const clickedEl = event.target;
if (clickedEl !== selectedAnnoteEl) {
unselectCodeLines();
const activeEl = window.document.querySelector('dt[data-target-cell].code-annotation-active');
if (activeEl) {
activeEl.classList.remove('code-annotation-active');
}
selectCodeLines(clickedEl);
clickedEl.classList.add('code-annotation-active');
} else {
// Unselect the line
unselectCodeLines();
clickedEl.classList.remove('code-annotation-active');
}
});
}
const findCites = (el) => {
const parentEl = el.parentElement;
if (parentEl) {
const cites = parentEl.dataset.cites;
if (cites) {
return {
el,
cites: cites.split(' ')
};
} else {
return findCites(el.parentElement)
}
} else {
return undefined;
}
};
var bibliorefs = window.document.querySelectorAll('a[role="doc-biblioref"]');
for (var i=0; i<bibliorefs.length; i++) {
const ref = bibliorefs[i];
const citeInfo = findCites(ref);
if (citeInfo) {
tippyHover(citeInfo.el, function() {
var popup = window.document.createElement('div');
citeInfo.cites.forEach(function(cite) {
var citeDiv = window.document.createElement('div');
citeDiv.classList.add('hanging-indent');
citeDiv.classList.add('csl-entry');
var biblioDiv = window.document.getElementById('ref-' + cite);
if (biblioDiv) {
citeDiv.innerHTML = biblioDiv.innerHTML;
}
popup.appendChild(citeDiv);
});
return popup.innerHTML;
});
}
}
});
</script>
</div> <!-- /content -->
</body></html>
Reference in New Issue View Git Blame Copy Permalink