Custom Integrations

Axolotl adds custom features through integrations. They are located within the src/axolotl/integrations directory.

To enable them, please check the respective documentations.

Cut Cross Entropy

Cut Cross Entropy (CCE) reduces VRAM usage through optimization on the cross-entropy operation during loss calculation.

See https://github.com/apple/ml-cross-entropy

Requirements

  • PyTorch 2.4.0 or higher

Installation

Run the following command to install cut_cross_entropy[transformers] if you don’t have it already.

  • If you are in dev environment
python scripts/cutcrossentropy_install.py | sh
  • If you are installing from pip
pip3 uninstall -y cut-cross-entropy && pip3 install "cut-cross-entropy[transformers] @ git+https://github.com/axolotl-ai-cloud/ml-cross-entropy.git@3574df5"

Usage

plugins:
  - axolotl.integrations.cut_cross_entropy.CutCrossEntropyPlugin

Supported Models

  • afmoe
  • apertus
  • arcee
  • cohere
  • cohere2
  • cohere2_moe
  • cohere2_vision
  • cohere_compass
  • cohere_compass_text
  • deepseek_v2
  • deepseek_v3
  • deepseek_v4
  • exaone4
  • exaone4_5
  • exaone_moe
  • gemma
  • gemma2
  • gemma3
  • gemma3_text
  • gemma3n
  • gemma3n_text
  • gemma4
  • gemma4_text
  • gemma4_unified
  • gemma4_unified_text
  • glm
  • glm4
  • glm4_moe
  • glm4_moe_lite
  • glm46v
  • glm4v
  • glm4v_moe
  • glm_image
  • glm_moe_dsa
  • gpt_oss
  • granite
  • granitemoe
  • granitemoehybrid
  • granitemoeshared
  • hunyuan_v1_dense
  • hunyuan_v1_moe
  • internvl
  • kimi_linear
  • lfm2
  • lfm2_moe
  • lfm2_vl
  • llama
  • llama4
  • llama4_text
  • llava
  • minimax
  • minimax_m2
  • ministral
  • ministral3
  • mistral
  • mistral3
  • mistral4
  • mixtral
  • mllama
  • muse_glimmer
  • nemotron_h
  • olmo
  • olmo2
  • olmo3
  • olmoe
  • phi
  • phi3
  • phi4_multimodal
  • qwen2
  • qwen2_5_vl
  • qwen2_moe
  • qwen2_vl
  • qwen3
  • qwen3_5
  • qwen3_5_text
  • qwen3_5_moe
  • qwen3_5_moe_text
  • qwen3_moe
  • qwen3_next
  • qwen3_vl
  • qwen3_vl_moe
  • qwen4_exp
  • qwen4_exp_text
  • seed_oss
  • smollm3
  • step3p5
  • step3p7
  • voxtral

Citation

@article{wijmans2024cut,
  author       = {Erik Wijmans and
                  Brody Huval and
                  Alexander Hertzberg and
                  Vladlen Koltun and
                  Philipp Kr\"ahenb\"uhl},
  title        = {Cut Your Losses in Large-Vocabulary Language Models},
  journal      = {arXiv},
  year         = {2024},
  url          = {https://arxiv.org/abs/2411.09009},
}

Please see reference here

DenseMixer

See DenseMixer

Simply add the following to your axolotl YAML config:

plugins:
  - axolotl.integrations.densemixer.DenseMixerPlugin

Please see reference here

Diffusion LM Training Plugin for Axolotl

This plugin enables diffusion language model training using an approach inspired by LLaDA (Large Language Diffusion Models) within Axolotl.

Overview

LLaDA is a diffusion-based approach to language model training that uses: - Random token masking during training instead of next-token prediction - Bidirectional attention to allow the model to attend to the full context - Importance weighting based on masking probabilities for stable training

This approach can lead to more robust language models with better understanding of bidirectional context.

Installation

The plugin is included with Axolotl. See our installation docs.

Quickstart

Train with an example config (Llama‑3.2 1B): - Pretrain: axolotl train examples/llama-3/diffusion/pretrain-1b.yaml - SFT: axolotl train examples/llama-3/diffusion/sft-1b.yaml

Basic Configuration

You can also modify your existing configs to enable / customize diffusion training.

Add the following to your Axolotl config:

plugins:
  - axolotl.integrations.diffusion.DiffusionPlugin

And, configure the nested diffusion block (defaults shown):

diffusion:
  noise_schedule: linear  # or "cosine"
  min_mask_ratio: 0.1
  max_mask_ratio: 0.9
  num_diffusion_steps: 128
  eps: 1e-3
  importance_weighting: true

  # Mask token (training auto-adds if missing, avoid pad/eos)
  mask_token_str: "<|diffusion_mask|>"
  # Or use an existing special token id (e.g., 128002 for Llama-3.x)
  # mask_token_id: 128002

  # Sample generation during training (optional)
  generate_samples: true
  generation_interval: 100
  num_generation_samples: 3
  generation_steps: 128
  generation_temperature: 0.0
  generation_max_length: 100

Supported Models

Any models that support 4D attention masks should work out of the box. If not, please create an issue or open a PR!

How It Works

Random Masking

During training, tokens are randomly masked: - Sample timestep t uniformly from [0, 1] - Calculate masking probability: p = (1 - eps) * t + eps - Randomly mask tokens with probability p

Diffusion Loss

Loss is computed only on masked tokens with (optional) importance weighting:

loss = sum(cross_entropy(pred, target) / p_mask) / total_tokens

Sample Generation

When diffusion.generate_samples: true, the plugin generates samples during training:

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...

Samples are logged to console and wandb (if enabled).

Inference

Diffusion inference is integrated into the standard Axolotl CLI. Use the same config you trained with and run:

axolotl inference path/to/your-config.yaml

Optionally, pass --gradio to use a simple web interface.

Interactive controls (prefix the prompt with commands): - :complete N → completion mode with N new masked tokens appended (default 64) - :mask R → random masking mode with target mask ratio R in [0.0, 1.0]

Example session:

================================================================================
Commands:
:complete N -> completion mode with N tokens (default 64)
:mask R     -> 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

Metrics and Monitoring

The plugin adds (or modifies) several metrics to track diffusion training:

  • train/loss: Weighted diffusion loss
  • train/accuracy: Accuracy on masked tokens
  • train/mask_ratio: Average fraction of tokens masked
  • train/num_masked_tokens: Number of tokens masked
  • train/avg_p_mask: Average masking probability
  • train/ce_loss: Unweighted cross-entropy loss
  • train/importance_weight_avg: Average importance weight

Limitations

  • No flash attention support
  • No RL training support

References

Please see reference here

Expert Parallelism Integration

Replaces the MoE dispatch/combine path with DeepEP’s fused kernels.

Requirements

Ampere (sm_80, A100) or Hopper (sm_90, H100), all-pairs NVLink.

Installation

Hopper (sm_90, H100), multi-node with NCCL 2.29+ (torch 2.11+) and OFED:

git clone --depth 1 https://github.com/deepseek-ai/DeepEP.git
cd DeepEP
TORCH_CUDA_ARCH_LIST=9.0 MAX_JOBS=16 uv pip install --no-build-isolation .
python -c "import deep_ep; print(deep_ep.Buffer)"

Hopper (sm_90, H100), single-node intranode-only (no OFED):

git clone https://github.com/deepseek-ai/DeepEP.git
cd DeepEP
git checkout v1.2.1
git apply <<'EOF'
--- a/setup.py
+++ b/setup.py
@@ -19,7 +19,10 @@ if __name__ == '__main__':
     disable_nvshmem = False
     nvshmem_dir = os.getenv('NVSHMEM_DIR', None)
     nvshmem_host_lib = 'libnvshmem_host.so'
-    if nvshmem_dir is None:
+    if int(os.getenv('DISABLE_NVSHMEM', '0')):
+        disable_nvshmem = True
+        nvshmem_dir = None
+    elif nvshmem_dir is None:
         try:
             nvshmem_dir = importlib.util.find_spec("nvidia.nvshmem").submodule_search_locations[0]
             nvshmem_host_lib = get_nvshmem_host_lib_name(nvshmem_dir)
EOF

git apply <<'EOF'
--- a/setup.py
+++ b/setup.py
@@ -71,7 +71,9 @@ if __name__ == '__main__':
         os.environ['TORCH_CUDA_ARCH_LIST'] = os.getenv('TORCH_CUDA_ARCH_LIST', '9.0')

         # CUDA 12 flags
-        nvcc_flags.extend(['-rdc=true', '--ptxas-options=--register-usage-level=10'])
+        nvcc_flags.append('--ptxas-options=--register-usage-level=10')
+        if not disable_nvshmem:
+            nvcc_flags.append('-rdc=true')

     # Disable LD/ST tricks, as some CUDA version does not support `.L1::no_allocate`
     if os.environ['TORCH_CUDA_ARCH_LIST'].strip() != '9.0':
EOF

DISABLE_NVSHMEM=1 TORCH_CUDA_ARCH_LIST=9.0 MAX_JOBS=16 \
  uv pip install --no-build-isolation .
python -c "import deep_ep; print(deep_ep.Buffer)"

Notes:

  • Hopper kernels (FP8, TMA, etc.) are preserved; only intranode dispatch/combine is built — appropriate for single-node H100×{4,8}.
  • Patch 1 lets DISABLE_NVSHMEM=1 skip the NVSHMEM build path, which would otherwise need Mellanox OFED dev headers (infiniband/mlx5dv.h).
  • Patch 2 drops -rdc=true when NVSHMEM is off; otherwise the device-link step has nothing to link against and import fails with __cudaRegisterLinkedBinary_* undefined symbol.
  • The v1.2.1 pin is required: it is the last release whose setup.py still carries the disable_nvshmem path the two patches edit. DeepEP main restructured setup.py and removed that path (DeepEP #664), so against HEAD the patches fail with patch failed: setup.py:19. The pin also sidesteps DeepEP HEAD’s csrc/elastic/ (Engram/EPv2, commit b306af0), which needs ncclGinRequest_t from NCCL 2.29+.

Ampere (sm_80, A100, intranode-only) — needs two small source patches gated on DISABLE_NVSHMEM=1:

git clone https://github.com/deepseek-ai/DeepEP.git
cd DeepEP
git checkout v1.2.1

git apply <<'EOF'
--- a/setup.py
+++ b/setup.py
@@ -19,7 +19,10 @@ if __name__ == '__main__':
     disable_nvshmem = False
     nvshmem_dir = os.getenv('NVSHMEM_DIR', None)
     nvshmem_host_lib = 'libnvshmem_host.so'
-    if nvshmem_dir is None:
+    if int(os.getenv('DISABLE_NVSHMEM', '0')):
+        disable_nvshmem = True
+        nvshmem_dir = None
+    elif nvshmem_dir is None:
         try:
             nvshmem_dir = importlib.util.find_spec("nvidia.nvshmem").submodule_search_locations[0]
             nvshmem_host_lib = get_nvshmem_host_lib_name(nvshmem_dir)
EOF

git apply <<'EOF'
--- a/csrc/deep_ep.cpp
+++ b/csrc/deep_ep.cpp
@@ -1823,22 +1823,34 @@ bool is_sm90_compiled() {
 }

 void Buffer::low_latency_update_mask_buffer(int rank_to_mask, bool mask) {
+#ifndef DISABLE_NVSHMEM
     EP_HOST_ASSERT(mask_buffer_ptr != nullptr and "Shrink mode must be enabled");
     EP_HOST_ASSERT(rank_to_mask >= 0 and rank_to_mask < num_ranks);
     internode_ll::update_mask_buffer(mask_buffer_ptr, rank_to_mask, mask, at::cuda::getCurrentCUDAStream());
+#else
+    EP_HOST_ASSERT(false and "NVSHMEM is disabled during compilation");
+#endif
 }

 void Buffer::low_latency_query_mask_buffer(const torch::Tensor& mask_status) {
+#ifndef DISABLE_NVSHMEM
     EP_HOST_ASSERT(mask_buffer_ptr != nullptr and "Shrink mode must be enabled");
     EP_HOST_ASSERT(mask_status.numel() == num_ranks && mask_status.scalar_type() == torch::kInt32);

     internode_ll::query_mask_buffer(
         mask_buffer_ptr, num_ranks, reinterpret_cast<int*>(mask_status.data_ptr()), at::cuda::getCurrentCUDAStream());
+#else
+    EP_HOST_ASSERT(false and "NVSHMEM is disabled during compilation");
+#endif
 }

 void Buffer::low_latency_clean_mask_buffer() {
+#ifndef DISABLE_NVSHMEM
     EP_HOST_ASSERT(mask_buffer_ptr != nullptr and "Shrink mode must be enabled");
     internode_ll::clean_mask_buffer(mask_buffer_ptr, num_ranks, at::cuda::getCurrentCUDAStream());
+#else
+    EP_HOST_ASSERT(false and "NVSHMEM is disabled during compilation");
+#endif
 }

 }  // namespace deep_ep
EOF

DISABLE_NVSHMEM=1 DISABLE_SM90_FEATURES=1 TORCH_CUDA_ARCH_LIST=8.0 MAX_JOBS=16 \
  uv pip install --no-build-isolation .
python -c "import deep_ep; print(deep_ep.Buffer)"

Usage

plugins:
  - axolotl.integrations.expert_parallel.ExpertParallelPlugin

expert_parallel_size: 2  # 1 = disabled (default); > 1 = enabled

For composition with FSDP at 4+ GPUs, set both expert_parallel_size and dp_shard_size. The product must equal world_size:

expert_parallel_size: 2
dp_shard_size: 2
fsdp_version: 2
fsdp_config:
  auto_wrap_policy: TRANSFORMER_BASED_WRAP
  transformer_layer_cls_to_wrap: Qwen3MoeDecoderLayer
  state_dict_type: FULL_STATE_DICT
  reshard_after_forward: true

See full example configs at examples/expert_parallel/.

Implementation notes

EP composes with the local-experts kernel you’ve already configured: ScatterMoE, SonicMoE, grouped_mm, or eager.

EP composes with FSDP on orthogonal mesh axes: experts are sharded across the ep axis, non-expert params across dp_shard. The two collectives run on disjoint process groups, so they don’t conflict. Layout follows Expert Parallelism with FSDP (tinkerings.dev) — “rows share weights, columns move tokens.”

Your existing config Local kernel under DeepEP
use_scattermoe: true ScatterMoE (Triton)
use_sonicmoe: true SonicMoE (bf16 experts)
experts_implementation: grouped_mm / batched_mm grouped_mm (transformers)
experts_implementation: eager eager Python loop
(unset) grouped_mm (default)

Limitations

  • Models’ modeling code must use @use_experts_implementation (canonical 3D gate_up_proj / down_proj). ModuleList as used in Mixtral is not supported.
  • num_experts must be divisible by expert_parallel_size.
  • Supported mesh axes: EP, EP × dp_shard, EP × cp, EP × cp × dp_shard (experts shard on ep, the sequence on cp, non-expert weights on dp_shard). EP × TP is not yet supported and raises NotImplementedError. EP × CP requires the model’s attention to be context-parallel-aware on the cp axis (e.g. GLM-5.2 DSA via the kernels plugin); stock attention uses accelerate CP.
  • DeepEP limitation: Low-latency (LL) kernels are inter-node only by design (pure RDMA via IBGDA). Single-node + intranode setups always use the standard kernels and don’t benefit from LL.
  • FP8 dispatch needs Hopper + DISABLE_SM90_FEATURES=0.

Troubleshooting

CUBLAS_STATUS_INVALID_VALUE on a basic bf16 GEMM after import deep_ep

The system’s libcublas.so.13 is older than what cu130 torch expects. Put the cu13 lib that ships with the torch wheel on LD_LIBRARY_PATH:

export LD_LIBRARY_PATH="$(python -c 'import nvidia.cu13 as m; print(list(m.__path__)[0] + "/lib")'):$LD_LIBRARY_PATH"

Unrelated to DeepEP itself, but anyone on cu130 torch hits it on boxes with a system CUDA toolkit older than 13.0.

CUDA error 803 (system not yet initialized)

On driver < 580, also prepend /usr/local/cuda-13.0/compat to LD_LIBRARY_PATH. Do not add the compat dir on driver ≥ 580 (its libcuda is older than the running driver and triggers CUDA error 803).

Please see reference here

Grokfast

See https://github.com/ironjr/grokfast

Usage

plugins:
  - axolotl.integrations.grokfast.GrokfastPlugin

grokfast_alpha: 2.0
grokfast_lamb: 0.98

Citation

@article{lee2024grokfast,
    title={{Grokfast}: Accelerated Grokking by Amplifying Slow Gradients},
    author={Lee, Jaerin and Kang, Bong Gyun and Kim, Kihoon and Lee, Kyoung Mu},
    journal={arXiv preprint arXiv:2405.20233},
    year={2024}
}

Please see reference here

Kernels Integration

MoE (Mixture of Experts) kernels speed up training for MoE layers and reduce VRAM costs. Transformers v5 introduced a uniform dispatch point for the per-expert grouped GEMMs via the experts_implementation config kwarg:

class ExpertsInterface(GeneralInterface):
    _global_mapping = {
        "batched_mm": batched_mm_experts_forward,
        "grouped_mm": grouped_mm_experts_forward,
        "sonicmoe":   sonicmoe_experts_forward,   # upstream HF integration
    }

Axolotl registers two additional implementations into this same global registry: ScatterMoE (Triton, runs on any CUDA GPU) and a LoRA-aware SonicMoE variant (CUTLASS / cute-DSL, Hopper or newer). Routing — softmax/sigmoid top-k, group selection, shared experts, bias correction, etc. — stays in each model’s SparseMoEBlock, where transformers handles all per-architecture variation. Axolotl only swaps the experts forward.

Usage

Add the following to your axolotl YAML config:

plugins:
  - axolotl.integrations.kernels.KernelsPlugin

use_kernels: true

use_scattermoe: true
use_sonicmoe: true

experts_implementation is auto-set to scattermoe / sonicmoe from the kernel flag, but you can override to eager / batched_mm / grouped_mm to compare against the transformers reference implementations.

SonicMoE installation

Prerequisites: - NVIDIA Hopper (H100/H200) or Blackwell (B200/GB200/B300) GPU - CUDA 12.9+ (13.0+ for B300) - PyTorch 2.7+ - For B300: Triton 3.6.x

The sonic-moe kernel ships through the HF kernels package. Transformers v5.8+ auto-fetches a prebuilt kernel from kernels-community/sonic-moe on first use:

uv pip install kernels "nvidia-cutlass-dsl==4.6.0" "apache-tvm-ffi>=0.1.10,<0.2"

apache-tvm-ffi is an undeclared runtime dependency of nvidia-cutlass-dsl 4.6.0 (absent from its Requires-Dist, so pip will not pull it); <0.1.10 breaks cute.compile, so pin it explicitly.

Note: Blackwell support is in upstream beta. On Blackwell GPUs Axolotl automatically sets USE_QUACK_GEMM=1 to enable the Blackwell kernels.

How It Works

The KernelsPlugin runs once before model loading and:

  1. Calls register_scattermoe_experts() or register_sonicmoe_experts(), which inserts the kernel forward into transformers.integrations.moe.ALL_EXPERTS_FUNCTIONS.
  2. Sets cfg.experts_implementation to the matching name.
  3. When the model loads, transformers’ @use_experts_implementation decorator on each model’s Experts class reads config._experts_implementation and dispatches to our registered forward.

That’s the entire integration — there is no per-architecture SparseMoEBlock monkey-patch, no per-model routing code, and no weight-layout conversion. As new MoE models adopt the decorator upstream they immediately benefit from both kernels.

BF16 LoRA Support

Both kernels train PEFT adapters on gate_up_proj / down_proj (and gate for the router) end-to-end:

  • ScatterMoE fuses the LoRA B @ A product into the per-expert grouped GEMM via custom Triton kernels (parallel_linear_lora). No extra materialization pass.
  • SonicMoE materializes W_eff = W + scaling * (B @ A) per expert inside a custom MoELoRAMaterialize autograd.Function and passes the effective weight into the CUTLASS kernel. Backward decomposes dW_eff into dA and dB via the chain rule, so LoRA parameters train without modifying the kernel.

Both paths detect PEFT ParamWrapper on individual expert parameters (target_parameters API) and unwrap them before dispatch.

ScatterMoE NVFP4 (W4A16) LoRA

Train LoRA on ModelOpt NVFP4 checkpoint via ScatterMoE. Routed experts are dequantized to bf16 (W4A16).

Requires: - CUDA GPU with Triton. - qwen3_moe, qwen3_next, deepseek_v4, glm_moe_dsa, and gemma4_text

Tip: in our tests, the Triton dequant path below is currently faster end to end than the ScatterMoE NVFP4 (if the arch is supported).

SonicMoE NVFP4 (W4A4) LoRA

Train LoRA on ModelOpt NVFP4 checkpoint via Quack (e.g. nvidia/Qwen3-30B-A3B-NVFP4).

Requires: - Blackwell SM100 for W4A4, others for W4A16 - qwen3_moe / qwen3_next

Install the pinned quack kernels (other versions untested):

uv pip install "quack-kernels==0.6.1" "nvidia-cutlass-dsl==4.6.0" "apache-tvm-ffi>=0.1.10,<0.2"

AXOLOTL_SONICMOE_NVFP4_BACKEND picks the expert GEMM (unset = auto):

Backend Compute Runs on
fp4_cute native W4A4 tensor cores (quack) Blackwell B200/GB200 (SM100/110)
dequant dequant to bf16 (W4A16) any CUDA GPU with Triton

Advanced tuning knobs (fused up-proj, per-tensor-scale fold, fp8 DeepGEMM backward) are exposed as other AXOLOTL_SONICMOE_NVFP4_* env vars; defaults are correct for normal training.

Tip: in our tests, the Triton dequant path is currently faster (and lower mem) end to end than the fp4_cute path for < 100B param MoE model. However, fp4_cute should overtake when expert matmul become more dominant cost.

Model Support

Any model whose Experts class is decorated with @use_experts_implementation upstream works automatically. As of transformers 5.8 this includes (verified). The bf16 columns are base (unquantized) support; the NVFP4 columns mark which kernel trains LoRA on a ModelOpt NVFP4 checkpoint of that arch:

Model Type ScatterMoE (bf16) SonicMoE (bf16) NVFP4 (ScatterMoE) NVFP4 (SonicMoE)
mixtral Yes Yes - -
qwen2_moe Yes Yes - -
qwen3_moe Yes Yes Yes Yes
qwen3_next Yes Yes Yes Yes
qwen3_5_moe Yes Yes - -
olmoe Yes Yes - -
mistral4 Yes Yes - -
glm_moe_dsa Yes Yes Yes -
deepseek_v3 Yes Yes - -
minimax_m2 Yes Yes - -
ernie4_5_moe Yes Yes - -
hunyuan_v1_moe Yes Yes - -
gemma4_text Yes Yes Yes -
gpt_oss Yes No - -
nemotron_h Yes Yes¹ Yes Yes¹

NVFP4 for deepseek_v4 is supported via ScatterMoE with use_dsv4_kernels (its own fused-kernel path), so it is not a row above.

¹ nemotron_h covers Nemotron-3 latentmoe: non-gated relu² experts (up_proj/down_proj, no gate), operating in moe_latent_size when set (tokens arrive pre-projected by the block’s fc1/fc2_latent_proj). ScatterMoE handles that layout natively; SonicMoE routes it through a torch._grouped_mm MLP because the current sonic-moe op layer only allows gated epilogues, and AXOLOTL_SONICMOE_NONGATED_FUSED=1 switches to the exact REGLU rewrite (relu²(h) == h·relu(h)) once a build relaxes that assert. Nemotron-3 NVFP4 checkpoints (modelopt MIXED_PRECISION) keep the routed experts packed NVFP4 and dequantize every other quantized linear to bf16 at load; quantization there is per layer, not per module name, so one converter per suffix picks its branch from the keys each layer ships.

gpt_oss carries the decorator with is_concatenated=False, is_transposed=True, has_bias=True and uses a clamped sigmoid-GLU activation. The ScatterMoE forward handles the transposed/interleaved/biased layout and that epilogue via its Triton path (no weight transpose, interleaved gate/up, per-expert bias folded into the grouped GEMM).

Epilogue check (why gpt_oss is No on SonicMoE)

SonicMoE picks its fused epilogue from config.hidden_act, which cannot express a clamped or otherwise non-plain GLU. GptOssConfig.hidden_act is "silu", so gpt_oss silently ran plain SwiGLU (10.4% top-1 agreement, teacher-forced NLL 2.30 -> 8.12 on openai/gpt-oss-20b).

Axolotl now probes each model’s declared _apply_gate numerically at load and raises if the chosen path cannot reproduce it, pointing at expert_backend: scattermoe.

Feature comparison

Feature ScatterMoE SonicMoE
Kernel backend Triton CUTLASS / cute-DSL
GPU requirement Any CUDA Hopper+
LoRA path Fused in Triton kernel MoELoRAMaterialize + custom autograd
LoRA overhead Lower (fused) Higher (materialization pass)
Selective expert dequantization Yes (~97% memory savings) No
Weight format Standard [E, 2*I, H] Standard [E, 2*I, H] (concat layout, no interleave)

Note on MegaBlocks

We tested MegaBlocks but were unable to ensure numerical accuracy, so we did not integrate it. It was also incompatible with many newer model architectures in transformers.

Please see reference here

Knowledge Distillation (KD)

Usage

plugins:
  - "axolotl.integrations.kd.KDPlugin"

kd_trainer: True
kd_ce_alpha: 0.1
kd_alpha: 0.9
kd_temperature: 1.0

torch_compile: True  # recommended to reduce vram

datasets:
  - path: ...
    type: "axolotl.integrations.kd.chat_template"
    field_messages: "messages_combined"
    logprobs_field: "llm_text_generation_vllm_logprobs"  # for kd only, field of logprobs

An example dataset can be found at axolotl-ai-co/evolkit-logprobs-pipeline-75k-v2-sample

Please see reference here

LLMCompressor

Fine-tune sparsified models in Axolotl using Neural Magic’s LLMCompressor.

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.

It uses Axolotl’s plugin system to hook into the fine-tuning flows while maintaining sparsity throughout training.

Requirements

  • Axolotl with llmcompressor extras:

    pip install "axolotl[llmcompressor]"

  • Requires llmcompressor >= 0.5.1

This will install all necessary dependencies to fine-tune sparsified models using the integration.

Usage

To enable sparse fine-tuning with this integration, include the plugin in your Axolotl config:

plugins:
  - axolotl.integrations.llm_compressor.LLMCompressorPlugin

llmcompressor:
  recipe:
    finetuning_stage:
      finetuning_modifiers:
        ConstantPruningModifier:
          targets: [
            're:.*q_proj.weight',
            're:.*k_proj.weight',
            're:.*v_proj.weight',
            're:.*o_proj.weight',
            're:.*gate_proj.weight',
            're:.*up_proj.weight',
            're:.*down_proj.weight',
          ]
          start: 0
  save_compressed: true

This plugin does not apply pruning or sparsification itself — it is intended for fine-tuning models that have already been sparsified.

Pre-sparsified checkpoints can be: - Generated using LLMCompressor - Downloaded from Neural Magic’s Hugging Face page - Any custom LLM with compatible sparsity patterns that you’ve created yourself

To learn more about writing and customizing LLMCompressor recipes, refer to the official documentation: https://github.com/vllm-project/llm-compressor/blob/main/README.md

Storage Optimization with save_compressed

Setting save_compressed: true 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)

This option is highly recommended when working with sparse models to maximize the benefits of model compression.

Example Config

See examples/llama-3/sparse-finetuning.yaml for a complete example.

Inference with vLLM

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.:

from vllm import LLM, SamplingParams

prompts = [
    "Hello, my name is",
    "The president of the United States is",
    "The capital of France is",
    "The future of AI is",
]
sampling_params = SamplingParams(temperature=0.8, top_p=0.95)
llm = LLM("path/to/your/sparse/model")
outputs = llm.generate(prompts, sampling_params)

for output in outputs:
    prompt = output.prompt
    generated_text = output.outputs[0].text
    print(f"Prompt: {prompt!r}, Generated text: {generated_text!r}")

For more details on vLLM’s capabilities and advanced configuration options, see the official vLLM documentation.

Learn More

For details on available sparsity and quantization schemes, fine-tuning recipes, and usage examples, visit the official LLMCompressor repository:

https://github.com/vllm-project/llm-compressor

Please see reference here

Language Model Evaluation Harness (LM Eval)

Run evaluation on model using the popular lm-evaluation-harness library.

See https://github.com/EleutherAI/lm-evaluation-harness

Usage

There are two ways to use the LM Eval integration:

1. Post-Training Evaluation

When training with the plugin enabled, evaluation runs automatically after training completes:

plugins:
  - axolotl.integrations.lm_eval.LMEvalPlugin

lm_eval_tasks:
  - gsm8k
  - hellaswag
  - arc_easy

lm_eval_batch_size: # Batch size for evaluation

output_dir:

Run training as usual:

axolotl train config.yml

2. Standalone CLI Evaluation

Evaluate any model directly without training:

lm_eval_model: meta-llama/Llama-2-7b-hf

plugins:
  - axolotl.integrations.lm_eval.LMEvalPlugin

lm_eval_tasks:
  - gsm8k
  - hellaswag
  - arc_easy

lm_eval_batch_size: 8
output_dir: ./outputs

Run evaluation:

axolotl lm-eval config.yml

Model Selection Priority

The model to evaluate is selected in the following priority order:

  1. lm_eval_model - Explicit model path or HuggingFace repo (highest priority)
  2. hub_model_id - Trained model pushed to HuggingFace Hub
  3. output_dir - Local checkpoint directory containing trained model weights

Citation

@misc{eval-harness,
  author       = {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},
  title        = {A framework for few-shot language model evaluation},
  month        = 07,
  year         = 2024,
  publisher    = {Zenodo},
  version      = {v0.4.3},
  doi          = {10.5281/zenodo.12608602},
  url          = {https://zenodo.org/records/12608602}
}

Please see reference here

Liger Kernels

Liger Kernel provides efficient Triton kernels for LLM training, offering:

  • 20% increase in multi-GPU training throughput
  • 60% reduction in memory usage
  • Compatibility with both FSDP and DeepSpeed

See https://github.com/linkedin/Liger-Kernel

Usage

plugins:
  - axolotl.integrations.liger.LigerPlugin
liger_rope: true
liger_rms_norm: true
liger_glu_activation: true
liger_layer_norm: true
liger_fused_linear_cross_entropy: true

liger_use_token_scaling: true

liger_kernel_impl: cutedsl

Supported Models

Any model type in liger-kernel’s native dispatch table (liger_kernel.transformers.monkey_patch.MODEL_TYPE_TO_APPLY_LIGER_FN) is supported out of the box — llama, mistral, mixtral, qwen2/qwen3 families, gemma through gemma4, deepseek_v4, glm4, phi3, olmo2, paligemma, and many more (50+ types as of liger 0.8.1).

On top of the native table, axolotl hand-patches these types (not covered upstream, or extended with kernels the native path lacks):

  • deepseek_v2
  • gemma4_unified / gemma4_unified_text
  • granitemoe
  • jamba
  • qwen3_5 / qwen3_5_moe (adds the fused gated-RMSNorm kernel for linear-attention layers)

Citation

@article{hsu2024ligerkernelefficienttriton,
      title={Liger Kernel: Efficient Triton Kernels for LLM Training},
      author={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},
      year={2024},
      eprint={2410.10989},
      archivePrefix={arXiv},
      primaryClass={cs.LG},
      url={https://arxiv.org/abs/2410.10989},
      journal={arXiv preprint arXiv:2410.10989},
}

Please see reference here

NeMo Gym Integration for Axolotl

Train LLMs with reinforcement learning using NVIDIA NeMo Gym environments as reward sources. NeMo Gym provides 50+ verified RL environments spanning math, coding, tool-use, reasoning, and safety — each with deterministic reward signals.

Validated Training Paths

Path Speed Multi-turn Architecture
Async GRPO + Data Producer Fastest (3x) Yes NemoGymDataProducer replaces vLLM generation
Standard GRPO + Data Producer Baseline Yes Same producer, no async prefetch
Standard GRPO + /verify Simplest No Reward function calls /verify directly
FSDP2 + /verify (2 GPU) Distributed No fsdp_version: 2

Multi-turn uses nemo_gym_multi_turn: true which auto-enables the async trainer’s data producer protocol. The plugin’s NemoGymDataProducer calls NeMo Gym agent /run endpoints and returns RolloutDataset with proper IS correction, env_mask, and rewards.

All paths tested end-to-end with Qwen3-0.6B + LoRA, logged to wandb project nemo-gym-rl.

Quick Start

Prerequisites

  • uv package manager (for NeMo Gym’s venv)
  • Two GPUs recommended (one for vLLM server, one for training)

1. Set Up NeMo Gym

git clone https://github.com/NVIDIA-NeMo/Gym.git ~/Gym
cd ~/Gym
uv venv --python 3.12 && source .venv/bin/activate && uv sync

CFLAGS="" uv pip install pycosat --python .venv/bin/python --no-build-isolation

for dir in resources_servers/reasoning_gym resources_servers/example_single_tool_call responses_api_models/vllm_model responses_api_agents/simple_agent; do
    uv venv --seed --allow-existing --python 3.12 $dir/.venv
    CFLAGS="" uv pip install --python $dir/.venv/bin/python pycosat --no-build-isolation 2>/dev/null
    uv pip install --python $dir/.venv/bin/python -e . "ray[default]==2.52.1"
done

uv pip install --python resources_servers/reasoning_gym/.venv/bin/python \
    reasoning-gym matplotlib pillow cycler contourpy kiwisolver

3. Single-Turn Training (Simplest — No Agent Server Needed)

For environments that only need single-turn verify (math, coding challenges), you don’t need an agent server. The plugin’s reward function calls /verify directly.

base_model: Qwen/Qwen2.5-0.5B-Instruct
rl: grpo
chat_template: tokenizer_default

trl:
  use_vllm: true
  vllm_mode: colocate
  vllm_enable_sleep_mode: false
  num_generations: 8
  max_completion_length: 128
  temperature: 0.9
  reward_funcs:
    - axolotl.integrations.nemo_gym.rewards.reward_nemo_gym_verify

plugins:
  - axolotl.integrations.nemo_gym.NemoGymPlugin

nemo_gym_enabled: true
nemo_gym_auto_start: false
nemo_gym_head_port: 11000
nemo_gym_datasets:
  - path: ~/Gym/resources_servers/reasoning_gym/data/train_basic_arithmetic.jsonl
    server_name: reasoning_gym

datasets:
  - path: ~/Gym/resources_servers/reasoning_gym/data/train_basic_arithmetic.jsonl
    type: chat_template
    field_messages: responses_create_params.input
    message_field_content: content
    message_field_role: role

vllm:
  gpu_memory_utilization: 0.3
  max_model_len: 512
  tensor_parallel_size: 1

learning_rate: 1e-5
micro_batch_size: 4
gradient_accumulation_steps: 2
max_steps: 50
output_dir: ./outputs/nemo_gym_arithmetic

Only needs ng_run with resource servers (no agent config):

cd ~/Gym && ng_run "+config_paths=[resources_servers/reasoning_gym/configs/resources_only.yaml]" "+skip_venv_if_present=true"

How It Works

Single-Turn

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

Multi-Turn (Agent /run)

┌─────────────┐     ┌──────────────┐     ┌──────────────────┐
│  axolotl    │     │  NeMo Gym    │────▶│  vLLM OpenAI     │
│  train      │────▶│  Agent /run  │◀────│  Server (GPU 0)  │
│  (GPU 1)    │     │              │     │  /v1/completions  │
└─────────────┘     └──────┬───────┘     └──────────────────┘
                           │
                           ▼
                    ┌──────────────┐
                    │  Resource    │
                    │  Server     │
                    │  (tools +   │
                    │   verify)   │
                    └─────────────┘

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

Data Producer Architecture (Multi-Turn)

When nemo_gym_multi_turn: true, the plugin automatically forces use_data_producer: true which selects the AxolotlAsyncGRPOTrainer. The plugin then swaps the trainer’s data producer with NemoGymDataProducer, which:

  1. Gets a prompt batch from the dataset iterator
  2. Expands by num_generations (one agent call per rollout)
  3. Calls NeMo Gym agents via async HTTP (aiohttp.gather)
  4. Parses responses into padded tensors (RolloutDataset)
  5. Returns with _pending_policy_logps=True for deferred scoring

The main thread then runs _compute_deferred_scores() which: - Computes policy logprobs on the training model (GPU forward pass) - Computes IS correction 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

With async_prefetch: true, the data producer runs in a background thread — giving ~3x speedup as generation and training overlap. With async_prefetch: false, it runs synchronously on the main thread (still uses the data producer protocol).

Weight Sync (LoRA Mode)

With vllm_lora_sync: true, the plugin (or async trainer) replaces NCCL-based weight sync with filesystem + HTTP:

  1. accelerator.get_state_dict() gathers LoRA weights from all ranks
  2. Rank 0 saves adapter to /tmp/lora_sync_*/vN/
  3. Rank 0 POSTs to /set_lora_adapter/ on vLLM server
  4. vLLM loads adapter natively via Punica kernels
  5. Only ~40MB transferred (vs multiple GBs for full model weights)

Multi-Environment Support

Datasets support per-row environment routing via agent_ref:

{"agent_ref": {"name": "reasoning_gym"}, "responses_create_params": {...}}
{"agent_ref": {"name": "instruction_following"}, "responses_create_params": {...}}

Or use the simpler per-dataset routing:

nemo_gym_datasets:
  - path: reasoning_data.jsonl
    server_name: reasoning_gym
  - path: tool_data.jsonl
    server_name: example_single_tool_call

Configuration Reference

Parameter Type Default Description
nemo_gym_enabled bool null Enable the NeMo Gym integration
nemo_gym_dir str ~/Gym Path to NeMo Gym repo
nemo_gym_auto_clone bool true Auto-clone NeMo Gym repo if missing
nemo_gym_auto_start bool true Auto-start resource servers
nemo_gym_config_paths list[str] — Server config YAMLs (relative to gym_dir)
nemo_gym_datasets list[dict] required Dataset configs with path and optional server_name
nemo_gym_head_port int 11000 Head server port
nemo_gym_server_timeout int 360 Server startup timeout (seconds)
nemo_gym_verify_timeout int 30 Per-request timeout (seconds)
nemo_gym_multi_turn bool false Enable multi-turn via agent /run

Dataset JSONL Format

Each line must have responses_create_params with input messages:

{
  "responses_create_params": {
    "input": [{"role": "user", "content": "What's the weather in SF?"}],
    "tools": [{"name": "get_weather", "type": "function", "strict": true, "parameters": {...}}]
  }
}

For multi-turn agent routing, include agent_ref:

{"agent_ref": {"name": "my_agent"}, "responses_create_params": {...}}

Note: Tool definitions MUST include "strict": true and "additionalProperties": false for NeMo Gym agent compatibility.

Reward Functions

The plugin provides two built-in reward functions — no user code needed:

trl:
  reward_funcs:
    # Multi-turn (nemo_gym_multi_turn: true):
    # Passthrough — agent /run already computed the reward
    - axolotl.integrations.nemo_gym.rewards.reward_env

    # Single-turn (nemo_gym_multi_turn: false):
    # Calls /verify endpoints on NeMo Gym resource servers
    - axolotl.integrations.nemo_gym.rewards.reward_nemo_gym_verify

Both are also importable from Python:

from axolotl.integrations.nemo_gym import reward_env, reward_nemo_gym_verify

Known Issues / Troubleshooting

NeMo Gym Server Setup

  • pycosat build failure: CFLAGS="" uv pip install pycosat --no-build-isolation
  • Ray version mismatch: Pin ray[default]==2.52.1 in all server venvs
  • Pre-build venvs: ng_run creates per-server venvs via Ray. Pre-build them and use +skip_venv_if_present=true
  • Tool strict field required: Agent server validates tool definitions require strict: true

vLLM / Weight Sync

  • Start vLLM with LoRA + tool calling + runtime loading:

    VLLM_ALLOW_RUNTIME_LORA_UPDATING=1 \
CUDA_VISIBLE_DEVICES=0 python -m vllm.entrypoints.openai.api_server \
  --model Qwen/Qwen3-4B-Instruct-2507 \
  --max-model-len 4096 \
  --gpu-memory-utilization 0.7 \
  --enable-lora --max-lora-rank 64 \
  --enable-auto-tool-choice --tool-call-parser hermes

  • VLLM_ALLOW_RUNTIME_LORA_UPDATING=1: Required for vllm_lora_sync: true. Without it, vLLM won’t expose the /v1/load_lora_adapter endpoint and weight sync will fail silently. The plugin warns if this endpoint is missing.

  • --enable-lora: Enables LoRA adapter support in vLLM

  • --enable-auto-tool-choice --tool-call-parser hermes: Required for Qwen3 tool calling

  • max_model_len must be > max_completion_length: Leave room for prompt tokens (~200). If equal, the NeMo Gym model proxy gets a 400 error and returns empty completions.

  • CUDA_HOME required: DeepSpeed import needs it for the nvcc shim

  • NCCL weight sync broken with vLLM 0.17: Use vllm_lora_sync: true (filesystem + HTTP via /v1/load_lora_adapter)

Multi-Turn

  • Agent server required: Multi-turn delegates to NeMo Gym’s agent server /run endpoint. Without an agent, the plugin falls back to single-turn /verify
  • Model server proxy: NeMo Gym needs a responses_api_models server that proxies to your vLLM. See the agent config example above

FSDP2

  • Validated on 2 GPUs with single-turn + LoRA
  • Async field filtering: The builder automatically filters async-only config fields when using the standard GRPO trainer

Comparison with Other Integrations

Feature Axolotl + NeMo Gym Unsloth + NeMo Gym NeMo RL (native)
Server management Automatic Manual (notebook) Built-in
Multi-environment Per-row routing Manual code YAML config
Multi-turn / tool use Agent /run delegation No Agent /run (Ray)
Async GRPO (3x speedup) Yes No Yes
LoRA sync Filesystem + HTTP N/A NCCL
Multi-GPU (FSDP2) Yes No Yes (Ray)
Config-driven Yes No (code) Yes

Please see reference here

Spectrum

by Eric Hartford, Lucas Atkins, Fernando Fernandes, David Golchinfar

This plugin contains code to freeze the bottom fraction of modules in a model, based on the Signal-to-Noise Ratio (SNR).

See https://github.com/cognitivecomputations/spectrum

Overview

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.

Usage

plugins:
  - axolotl.integrations.spectrum.SpectrumPlugin

spectrum_top_fraction: 0.5
spectrum_model_name: meta-llama/Meta-Llama-3.1-8B

Citation

@misc{hartford2024spectrumtargetedtrainingsignal,
      title={Spectrum: Targeted Training on Signal to Noise Ratio},
      author={Eric Hartford and Lucas Atkins and Fernando Fernandes Neto and David Golchinfar},
      year={2024},
      eprint={2406.06623},
      archivePrefix={arXiv},
      primaryClass={cs.LG},
      url={https://arxiv.org/abs/2406.06623},
}

Please see reference here

SwanLab Integration for Axolotl

SwanLab is an open-source, lightweight AI experiment tracking and visualization tool that provides a platform for tracking, recording, comparing, and collaborating on experiments.

This integration enables seamless experiment tracking and visualization of Axolotl training runs using SwanLab.

Features

  • 📊 Automatic Metrics Logging: Training loss, learning rate, and other metrics are automatically logged
  • 🎯 Hyperparameter Tracking: Model configuration and training parameters are tracked
  • 📈 Real-time Visualization: Monitor training progress in real-time through SwanLab dashboard
  • ☁️ Cloud & Local Support: Works in both cloud-synced and offline modes
  • 🔄 Experiment Comparison: Compare multiple training runs easily
  • 🤝 Team Collaboration: Share experiments with team members
  • 🎭 RLHF Completion Logging: Automatically log model outputs during DPO/KTO/ORPO/GRPO training for qualitative analysis
  • ⚡ Performance Profiling: Built-in profiling decorators to measure and optimize training performance
  • 🔔 Lark Notifications: Send real-time training updates to team chat (Feishu/Lark integration)

Installation

pip install swanlab

Quick Start

1. Register for SwanLab (Optional for cloud mode)

If you want to use cloud sync features, register at https://swanlab.cn to get your API key.

2. Configure Axolotl Config File

Add SwanLab configuration to your Axolotl YAML config:

plugins:
  - axolotl.integrations.swanlab.SwanLabPlugin

use_swanlab: true
swanlab_project: my-llm-project
swanlab_experiment_name: qwen-finetune-v1
swanlab_mode: cloud  # Options: cloud, local, offline, disabled
swanlab_workspace: my-team  # Optional: organization name
swanlab_api_key: YOUR_API_KEY  # Optional: can also use env var SWANLAB_API_KEY

3. Run Training

export SWANLAB_API_KEY=your-api-key-here

swanlab login

accelerate launch -m axolotl.cli.train your-config.yaml

Configuration Options

Basic Configuration

Parameter Type Default Description
use_swanlab bool false Enable SwanLab tracking
swanlab_project str None Project name (required)
swanlab_experiment_name str None Experiment name
swanlab_description str None Experiment description
swanlab_mode str cloud Sync mode: cloud, local, offline, disabled

Advanced Configuration

Parameter Type Default Description
swanlab_workspace str None Workspace/organization name
swanlab_api_key str None API key (prefer env var)
swanlab_web_host str None Private deployment web host
swanlab_api_host str None Private deployment API host
swanlab_log_model bool false Log model checkpoints (coming soon)
swanlab_lark_webhook_url str None Lark (Feishu) webhook URL for team notifications
swanlab_lark_secret str None Lark webhook HMAC secret for authentication
swanlab_log_completions bool true Enable RLHF completion table logging (DPO/KTO/ORPO/GRPO)
swanlab_completion_log_interval int 100 Steps between completion logging
swanlab_completion_max_buffer int 128 Max completions to buffer (memory bound)

Configuration Examples

Example 1: Basic Cloud Sync

plugins:
  - axolotl.integrations.swanlab.SwanLabPlugin

use_swanlab: true
swanlab_project: llama-finetune
swanlab_experiment_name: llama-3-8b-instruct-v1
swanlab_mode: cloud

Example 2: Offline/Local Mode

plugins:
  - axolotl.integrations.swanlab.SwanLabPlugin

use_swanlab: true
swanlab_project: local-experiments
swanlab_experiment_name: test-run-1
swanlab_mode: local  # or 'offline'

Example 3: Team Workspace

plugins:
  - axolotl.integrations.swanlab.SwanLabPlugin

use_swanlab: true
swanlab_project: research-project
swanlab_experiment_name: experiment-42
swanlab_workspace: my-research-team
swanlab_mode: cloud

Example 4: Private Deployment

plugins:
  - axolotl.integrations.swanlab.SwanLabPlugin

use_swanlab: true
swanlab_project: internal-project
swanlab_experiment_name: secure-training
swanlab_mode: cloud
swanlab_web_host: https://swanlab.yourcompany.com
swanlab_api_host: https://api.swanlab.yourcompany.com

Team Notifications with Lark (Feishu)

SwanLab supports sending real-time training notifications to your team chat via Lark (Feishu), ByteDance’s enterprise collaboration platform. This is especially useful for: - Production training monitoring: Get alerts when training starts, completes, or encounters errors - Team collaboration: Keep your ML team informed about long-running experiments - Multi-timezone teams: Team members can check training progress without being online

Prerequisites

  1. Lark Bot Setup: Create a custom bot in your Lark group chat
  2. Webhook URL: Get the webhook URL from your Lark bot settings
  3. HMAC Secret (recommended): Enable signature verification in your Lark bot for security

For detailed Lark bot setup instructions, see Lark Custom Bot Documentation.

Example 5: Basic Lark Notifications

Send training notifications to a Lark group chat:

plugins:
  - axolotl.integrations.swanlab.SwanLabPlugin

use_swanlab: true
swanlab_project: production-training
swanlab_experiment_name: llama-3-finetune-v2
swanlab_mode: cloud

swanlab_lark_webhook_url: https://open.feishu.cn/open-apis/bot/v2/hook/xxxxxxxxxx

Note: This configuration will work, but you’ll see a security warning recommending HMAC secret configuration.

Example 7: Team Workspace + Lark Notifications

Combine team workspace collaboration with Lark notifications:

plugins:
  - axolotl.integrations.swanlab.SwanLabPlugin

use_swanlab: true
swanlab_project: research-project
swanlab_experiment_name: multimodal-experiment-42
swanlab_workspace: ml-research-team
swanlab_mode: cloud

swanlab_lark_webhook_url: https://open.feishu.cn/open-apis/bot/v2/hook/xxxxxxxxxx
swanlab_lark_secret: your-webhook-secret-key

What Notifications Are Sent?

SwanLab’s Lark integration sends notifications for key training events: - Training Start: When your experiment begins - Training Complete: When training finishes successfully - Training Errors: If training crashes or encounters critical errors - Metric Milestones: Configurable alerts for metric thresholds (if configured in SwanLab)

Each notification includes: - Experiment name and project - Training status - Key metrics (loss, learning rate) - Direct link to SwanLab dashboard

Lark Configuration Validation

The plugin validates your Lark configuration at startup:

✅ Valid Configurations

use_swanlab: true
swanlab_project: my-project

use_swanlab: true
swanlab_project: my-project
swanlab_lark_webhook_url: https://open.feishu.cn/open-apis/bot/v2/hook/xxx
swanlab_lark_secret: your-secret

use_swanlab: true
swanlab_project: my-project
swanlab_lark_webhook_url: https://open.feishu.cn/open-apis/bot/v2/hook/xxx

Security Best Practices

  1. Always use HMAC secret in production:

    swanlab_lark_webhook_url: https://open.feishu.cn/...
swanlab_lark_secret: your-secret-key  # ✅ Add this!

  2. Store secrets in environment variables (even better):

    # In your training script/environment
export SWANLAB_LARK_WEBHOOK_URL="https://open.feishu.cn/..."
export SWANLAB_LARK_SECRET="your-secret-key"

    Then in config:

    # SwanLab plugin will auto-detect environment variables
use_swanlab: true
swanlab_project: my-project
# Lark URL and secret read from env vars

  3. Rotate webhook secrets periodically: Update your Lark bot’s secret every 90 days

  4. Use separate webhooks for dev/prod: Don’t mix development and production notifications

Distributed Training

Lark notifications are automatically deduplicated in distributed training: - Only rank 0 sends notifications - Other GPU ranks skip Lark registration - Prevents duplicate messages in multi-GPU training

torchrun --nproc_per_node=4 -m axolotl.cli.train config.yml

RLHF Completion Table Logging

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:

  • Inspect model behavior: See actual model outputs during training
  • Debug preference learning: Compare chosen vs rejected responses
  • Track reward patterns: Monitor how rewards evolve over training
  • Share examples with team: Visual tables in SwanLab dashboard

Features

  • ✅ Automatic detection: Works with DPO, KTO, ORPO, GRPO trainers
  • ✅ Memory-safe buffering: Bounded buffer prevents memory leaks in long training runs
  • ✅ Periodic logging: Configurable logging interval to reduce overhead
  • ✅ Rich visualization: SwanLab tables show prompts, responses, and metrics side-by-side

Configuration

Parameter Type Default Description
swanlab_log_completions bool true Enable completion logging for RLHF trainers
swanlab_completion_log_interval int 100 Log completions to SwanLab every N training steps
swanlab_completion_max_buffer int 128 Maximum completions to buffer (memory bound)

Example: DPO Training with Completion Logging

plugins:
  - axolotl.integrations.swanlab.SwanLabPlugin

use_swanlab: true
swanlab_project: dpo-training
swanlab_experiment_name: llama-3-dpo-v1
swanlab_mode: cloud

swanlab_log_completions: true
swanlab_completion_log_interval: 100  # Log every 100 steps
swanlab_completion_max_buffer: 128    # Keep last 128 completions

rl: dpo
datasets:
  - path: /path/to/preference_dataset
    type: chatml.intel

Example: Disable Completion Logging

If you’re doing a quick test run or don’t need completion tables:

plugins:
  - axolotl.integrations.swanlab.SwanLabPlugin

use_swanlab: true
swanlab_project: dpo-training

swanlab_log_completions: false

Supported RLHF Trainers

The completion logging callback automatically activates for these trainer types:

  • DPO (Direct Preference Optimization): Logs prompts, chosen, rejected, reward_diff
  • KTO (Kahneman-Tversky Optimization): Logs prompts, completions, labels, rewards
  • ORPO (Odds Ratio Preference Optimization): Logs prompts, chosen, rejected, log_odds_ratio
  • GRPO (Group Relative Policy Optimization): Logs prompts, completions, rewards, advantages
  • CPO (Constrained Policy Optimization): Logs prompts, chosen, rejected

For non-RLHF trainers (standard supervised fine-tuning), the completion callback is automatically skipped.

How It Works

  1. Auto-detection: Plugin detects trainer type at initialization
  2. Buffering: Completions are buffered in memory (up to swanlab_completion_max_buffer)
  3. Periodic logging: Every swanlab_completion_log_interval steps, buffer is logged to SwanLab
  4. Memory safety: Old completions are automatically dropped when buffer is full (uses collections.deque)
  5. Final flush: Remaining completions are logged when training completes

Viewing Completion Tables

After training starts, you can view completion tables in your SwanLab dashboard:

  1. Navigate to your experiment in SwanLab
  2. Look for the “rlhf_completions” table in the metrics panel
  3. The table shows:
    • step: Training step when completion was generated
    • prompt: Input prompt
    • chosen: Preferred response (DPO/ORPO)
    • rejected: Non-preferred response (DPO/ORPO)
    • completion: Model output (KTO/GRPO)
    • reward_diff/reward: Reward metrics
    • Trainer-specific metrics (e.g., log_odds_ratio for ORPO)

Memory Management

The completion buffer is memory-bounded to prevent memory leaks:

from collections import deque

buffer = deque(maxlen=128)  # Old completions automatically dropped

Memory usage estimate: - Average completion: ~500 characters (prompt + responses) - Buffer size 128: ~64 KB (negligible) - Buffer size 1024: ~512 KB (still small)

Recommendation: Default buffer size (128) works well for most cases. Increase to 512-1024 only if you need to review more historical completions.

Performance Impact

Completion logging has minimal overhead:

  • Buffering: O(1) append operation, negligible CPU/memory
  • Logging: Only happens every N steps (default: 100)
  • Network: SwanLab batches table uploads efficiently

Expected overhead: < 0.5% per training step

Troubleshooting

Completions not appearing in SwanLab

Cause: Trainer may not be logging completion data in the expected format.

Diagnostic steps: 1. Check trainer type detection in logs: text INFO: SwanLab RLHF completion logging enabled for DPOTrainer (type: dpo) 2. Verify your trainer is an RLHF trainer (DPO/KTO/ORPO/GRPO) 3. Check if trainer logs completion data (this depends on TRL version)

Note: The current implementation expects trainers to log completion data in the logs dict during on_log() callback. Some TRL trainers may not expose this data by default. You may need to patch the trainer to expose completions.

Buffer fills up too quickly

Cause: High logging frequency with small buffer size.

Solution: Increase buffer size or logging interval:

swanlab_completion_log_interval: 200  # Log less frequently
swanlab_completion_max_buffer: 512    # Larger buffer

Memory usage growing over time

Cause: Buffer should be bounded, so this indicates a bug.

Solution: 1. Verify swanlab_completion_max_buffer is set 2. Check SwanLab version is up to date 3. Report issue with memory profiling data

Performance Profiling

SwanLab integration includes profiling utilities to measure and log execution time of trainer methods. This helps you:

  • Identify bottlenecks: Find slow operations in your training loop
  • Optimize performance: Track improvements after optimization changes
  • Monitor distributed training: See per-rank timing differences
  • Debug hangs: Detect methods that take unexpectedly long

Features

  • ✅ Zero-config profiling: Automatic timing of key trainer methods
  • ✅ Decorator-based: Easy to add profiling to custom methods with @swanlab_profile
  • ✅ Context manager: Fine-grained profiling with swanlab_profiling_context()
  • ✅ Advanced filtering: ProfilingConfig for throttling and minimum duration thresholds
  • ✅ Exception-safe: Logs duration even if function raises an exception

Basic Usage: Decorator

Add profiling to any trainer method with the @swanlab_profile decorator:

from axolotl.integrations.swanlab.profiling import swanlab_profile

class MyCustomTrainer(AxolotlTrainer):
    @swanlab_profile
    def training_step(self, model, inputs):
        # Your training step logic
        return super().training_step(model, inputs)

    @swanlab_profile
    def prediction_step(self, model, inputs, prediction_loss_only):
        # Your prediction logic
        return super().prediction_step(model, inputs, prediction_loss_only)

The decorator automatically: 1. Measures execution time with high-precision timer 2. Logs to SwanLab as profiling/Time taken: ClassName.method_name 3. Only logs if SwanLab is enabled (use_swanlab: true) 4. Gracefully handles exceptions (logs duration, then re-raises)

Advanced Usage: Context Manager

For fine-grained profiling within a method:

from axolotl.integrations.swanlab.profiling import swanlab_profiling_context

class MyTrainer(AxolotlTrainer):
    def complex_training_step(self, model, inputs):
        # Profile just the forward pass
        with swanlab_profiling_context(self, "forward_pass"):
            outputs = model(**inputs)

        # Profile just the backward pass
        with swanlab_profiling_context(self, "backward_pass"):
            loss = outputs.loss
            loss.backward()

        return outputs

Advanced Usage: ProfilingConfig

Filter and throttle profiling logs with ProfilingConfig:

from axolotl.integrations.swanlab.profiling import (
    swanlab_profiling_context_advanced,
    ProfilingConfig,
)

profiling_config = ProfilingConfig(
    enabled=True,
    min_duration_ms=1.0,    # Only log if duration > 1ms
    log_interval=10,        # Log every 10th call
)

class MyTrainer(AxolotlTrainer):
    def frequently_called_method(self, data):
        with swanlab_profiling_context_advanced(
            self,
            "frequent_op",
            config=profiling_config
        ):
            # This only logs every 10th call, and only if it takes > 1ms
            result = expensive_computation(data)
        return result

ProfilingConfig Parameters: - enabled: Enable/disable profiling globally (default: True) - min_duration_ms: Minimum duration to log in milliseconds (default: 0.1) - log_interval: Log every Nth function call (default: 1 = log all)

Use cases: - High-frequency methods: Use log_interval=100 to reduce logging overhead - Filter noise: Use min_duration_ms=1.0 to skip very fast operations - Debugging: Use log_interval=1, min_duration_ms=0.0 to log everything

Viewing Profiling Metrics

In your SwanLab dashboard, profiling metrics appear under the “profiling” namespace:

profiling/Time taken: AxolotlTrainer.training_step
profiling/Time taken: AxolotlTrainer.prediction_step
profiling/Time taken: MyTrainer.forward_pass
profiling/Time taken: MyTrainer.backward_pass

You can: - Track over time: See if methods get faster/slower during training - Compare runs: Compare profiling metrics across experiments - Identify regressions: Detect if a code change slowed down training

Configuration in Axolotl Config

Profiling is automatically enabled when SwanLab is enabled. No additional config needed:

plugins:
  - axolotl.integrations.swanlab.SwanLabPlugin

use_swanlab: true
swanlab_project: my-project

To disable profiling while keeping SwanLab enabled:

from axolotl.integrations.swanlab.profiling import DEFAULT_PROFILING_CONFIG

DEFAULT_PROFILING_CONFIG.enabled = False

Performance Impact

  • Decorator overhead: ~2-5 microseconds per call (negligible)
  • Context manager overhead: ~1-3 microseconds (negligible)
  • Logging overhead: Only when SwanLab is enabled and method duration exceeds threshold
  • Network overhead: SwanLab batches metrics efficiently

Expected overhead: < 0.1% per training step (effectively zero)

Best Practices

  1. Profile bottlenecks first: Start by profiling suspected slow operations
  2. Use min_duration_ms: Filter out fast operations (< 1ms) to reduce noise
  3. Throttle high-frequency calls: Use log_interval for methods called > 100 times/step
  4. Profile across runs: Compare profiling metrics before/after optimization
  5. Monitor distributed training: Check for rank-specific slowdowns

Example: Complete Profiling Setup

from axolotl.integrations.swanlab.profiling import (
    swanlab_profile,
    swanlab_profiling_context,
    ProfilingConfig,
)

class OptimizedTrainer(AxolotlTrainer):
    def __init__(self, *args, **kwargs):
        super().__init__(*args, **kwargs)

        # Custom profiling config for high-frequency operations
        self.fast_op_config = ProfilingConfig(
            enabled=True,
            min_duration_ms=0.5,
            log_interval=50,
        )

    @swanlab_profile
    def training_step(self, model, inputs):
        """Main training step - always profile."""
        return super().training_step(model, inputs)

    @swanlab_profile
    def compute_loss(self, model, inputs, return_outputs=False):
        """Loss computation - always profile."""
        return super().compute_loss(model, inputs, return_outputs)

    def _prepare_inputs(self, inputs):
        """High-frequency operation - throttled profiling."""
        with swanlab_profiling_context_advanced(
            self,
            "prepare_inputs",
            config=self.fast_op_config,
        ):
            return super()._prepare_inputs(inputs)

Troubleshooting

Profiling metrics not appearing in SwanLab

Cause: SwanLab is not enabled or not initialized.

Solution:

use_swanlab: true
swanlab_project: my-project

Check logs for:

INFO: SwanLab initialized for project: my-project

Too many profiling metrics cluttering dashboard

Cause: Profiling every function call for high-frequency operations.

Solution: Use ProfilingConfig with throttling:

config = ProfilingConfig(
    min_duration_ms=1.0,    # Skip fast ops
    log_interval=100,       # Log every 100th call
)

Profiling overhead impacting training speed

Cause: Profiling itself should have negligible overhead (< 0.1%). If you see > 1% slowdown, this indicates a bug.

Solution: 1. Disable profiling temporarily to confirm: python DEFAULT_PROFILING_CONFIG.enabled = False 2. Report issue with profiling data and trainer details

Profiling shows inconsistent timing

Cause: Normal variation due to GPU warmup, data loading, or system load.

Solution: - Ignore first few steps (warmup period) - Look at average/median timing over many steps - Use log_interval to reduce noise from individual outliers

Complete Config Example

Here’s a complete example integrating SwanLab with your RVQ-Alpha training:

base_model: /path/to/your/model
model_type: Qwen2ForCausalLM

plugins:
  - axolotl.integrations.swanlab.SwanLabPlugin
  - axolotl.integrations.cut_cross_entropy.CutCrossEntropyPlugin

use_swanlab: true
swanlab_project: RVQ-Alpha-Training
swanlab_experiment_name: Qwen2.5-7B-MetaQA-Perturb-P020
swanlab_description: "Training on MetaQA and Perturbation datasets with NEW-RVQ encoding"
swanlab_mode: cloud
swanlab_workspace: single-cell-genomics

sequence_len: 32768
micro_batch_size: 1
gradient_accumulation_steps: 1
num_epochs: 2
learning_rate: 2e-5
optimizer: adamw_torch_fused

datasets:
  - path: /path/to/dataset
    type: chat_template

output_dir: ./outputs

Modes Explained

cloud Mode (Default)

  • Syncs experiments to SwanLab cloud in real-time
  • Requires API key and internet connection
  • Best for: Team collaboration, remote monitoring

local Mode

  • Saves experiments locally only
  • No cloud sync
  • Best for: Local development, air-gapped environments

offline Mode

  • Saves metadata locally
  • Can sync to cloud later using swanlab sync
  • Best for: Unstable internet, sync later

disabled Mode

  • Turns off SwanLab completely
  • No logging or tracking
  • Best for: Debugging, testing

Configuration Validation & Conflict Detection

SwanLab integration includes comprehensive validation and conflict detection to help you catch configuration errors early and avoid performance issues.

Required Fields Validation

The plugin validates your configuration at startup and provides clear error messages with solutions:

Missing Project Name

use_swanlab: true

Solution:

use_swanlab: true
swanlab_project: my-project

Invalid Mode

use_swanlab: true
swanlab_project: my-project
swanlab_mode: invalid-mode

Solution:

use_swanlab: true
swanlab_project: my-project
swanlab_mode: cloud  # or: local, offline, disabled

Empty Project Name

use_swanlab: true
swanlab_project: ""

Solution:

use_swanlab: true
swanlab_project: my-project

Cloud Mode API Key Warning

When using cloud mode without an API key, you’ll receive a warning with multiple solutions:

use_swanlab: true
swanlab_project: my-project
swanlab_mode: cloud

Solutions: 1. Set environment variable: export SWANLAB_API_KEY=your-api-key 2. Add to config (less secure): swanlab_api_key: your-api-key 3. Run swanlab login before training 4. Use swanlab_mode: local for offline tracking

Multi-Logger Performance Warnings

Using multiple logging tools simultaneously (SwanLab + WandB + MLflow + Comet) can impact training performance:

Two Loggers - Warning

use_swanlab: true
swanlab_project: my-project

use_wandb: true
wandb_project: my-project

Impact: - Performance overhead: ~1-2% per logger (cumulative) - Increased memory usage - Longer training time per step - Potential config/callback conflicts

Recommendations: - 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

Three+ Loggers - Error-Level Warning

use_swanlab: true
swanlab_project: my-project

use_wandb: true
wandb_project: my-project

use_mlflow: true
mlflow_tracking_uri: http://localhost:5000

Why This Matters: - 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

Auto-Enable Logic

For convenience, SwanLab will auto-enable if you specify a project without setting use_swanlab:

swanlab_project: my-project

use_swanlab: true
swanlab_project: my-project

Distributed Training Detection

In distributed training scenarios (multi-GPU), the plugin automatically detects and reports:

use_swanlab: true
swanlab_project: my-project
swanlab_mode: cloud

Why Only Rank 0: - Avoids duplicate experiment runs - Reduces network/cloud API overhead on worker ranks - Prevents race conditions in metric logging

Authentication

Method 2: Login Command

swanlab login

Method 3: Config File

swanlab_api_key: your-api-key-here

What Gets Logged?

Automatically Logged Metrics

  • Training loss
  • Learning rate
  • Gradient norm
  • Training steps
  • Epoch progress

Automatically Logged Config

  • Model configuration (base_model, model_type)
  • Training hyperparameters (learning_rate, batch_size, etc.)
  • Optimizer settings
  • Parallelization settings (FSDP, DeepSpeed, Context Parallel)
  • Axolotl configuration file
  • DeepSpeed configuration (if used)

Viewing Your Experiments

Cloud Mode

Visit https://swanlab.cn and navigate to your project to view: - Real-time training metrics - Hyperparameter comparison - System resource usage - Configuration files

Local Mode

swanlab watch ./swanlog

Integration with Existing Tools

SwanLab can work alongside other tracking tools:

plugins:
  - axolotl.integrations.swanlab.SwanLabPlugin

use_swanlab: true
swanlab_project: my-project

use_wandb: true
wandb_project: my-project

Troubleshooting

Configuration Errors

Error: “SwanLab enabled but ‘swanlab_project’ is not set”

Cause: You enabled SwanLab (use_swanlab: true) but forgot to specify a project name.

Solution:

use_swanlab: true
swanlab_project: my-project  # Add this line

Error: “Invalid swanlab_mode: ‘xxx’”

Cause: You provided an invalid mode value.

Solution: Use one of the valid modes:

swanlab_mode: cloud     # or: local, offline, disabled

Error: “swanlab_project cannot be an empty string”

Cause: You set swanlab_project: "" (empty string).

Solution: Either provide a valid name or remove the field:

swanlab_project: my-project

Import Errors

Error: “SwanLab is not installed”

Cause: SwanLab package is not installed in your environment.

Solution:

pip install swanlab
pip install swanlab>=0.3.0

Performance Issues

Warning: “Multiple logging tools enabled”

Cause: You have multiple experiment tracking tools enabled (e.g., SwanLab + WandB + MLflow).

Impact: ~1-2% performance overhead per logger, cumulative.

Solution: For production training, disable all but one logger:

use_swanlab: true
swanlab_project: my-project
use_wandb: false      # Disable others
use_mlflow: false

use_swanlab: false
use_wandb: true
wandb_project: my-project

Exception: Multiple loggers are acceptable for: - Short comparison runs (< 100 steps) - Migration testing between logging tools - Debugging logger-specific issues

Distributed Training Issues

SwanLab creates duplicate runs in multi-GPU training

Cause: All ranks are initializing SwanLab instead of just rank 0.

Expected Behavior: The plugin automatically ensures only rank 0 initializes SwanLab. You should see:

Info: Distributed training detected (world_size=4)
Info: Only rank 0 will initialize SwanLab
Info: Other ranks will skip SwanLab to avoid conflicts

If you see duplicates: 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

SwanLab not logging metrics

Solution: Ensure SwanLab is initialized before training starts. The plugin automatically handles this in pre_model_load.

API Key errors

Solution:

echo $SWANLAB_API_KEY

swanlab login

Cloud sync issues

Solution: Use offline mode and sync later:

swanlab_mode: offline

Then sync when ready:

swanlab sync ./swanlog

Plugin not loaded

Solution: Verify plugin path in config:

plugins:
  - axolotl.integrations.swanlab.SwanLabPlugin  # Correct path

Lark Notification Issues

Error: “Failed to import SwanLab Lark plugin”

Cause: Your SwanLab version doesn’t include the Lark plugin (requires SwanLab >= 0.3.0).

Solution:

pip install --upgrade swanlab

pip install 'swanlab>=0.3.0'

Warning: “Lark webhook has no secret configured”

Cause: You provided swanlab_lark_webhook_url but no swanlab_lark_secret.

Impact: Lark notifications will work, but without HMAC authentication (security risk).

Solution: Add HMAC secret for production use:

swanlab_lark_webhook_url: https://open.feishu.cn/open-apis/bot/v2/hook/xxx
swanlab_lark_secret: your-webhook-secret  # Add this line

When it’s OK to skip secret: - Local development and testing - Internal networks with restricted access - Non-sensitive training experiments

When secret is required: - Production training jobs - Training with proprietary data - Multi-team shared Lark groups

Error: “Failed to register Lark callback”

Cause: Invalid webhook URL or network connectivity issues.

Diagnostic steps:

curl -X POST "YOUR_WEBHOOK_URL" \
  -H 'Content-Type: application/json' \
  -d '{"msg_type":"text","content":{"text":"Test from Axolotl"}}'

pip show swanlab

Solution: 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

Lark notifications not received

Cause: Multiple possible causes.

Diagnostic checklist:

  1. Check training logs for Lark registration confirmation:

    # Expected log message (rank 0 only):
INFO: Registered Lark notification callback with HMAC authentication

  2. Verify webhook in Lark: Test webhook manually (see above)

  3. Check distributed training: Only rank 0 sends notifications

    # If running multi-GPU, check rank 0 logs specifically
grep "Registered Lark" logs/rank_0.log

  4. Verify SwanLab is initialized: Lark callback needs SwanLab to be running

    use_swanlab: true  # Must be enabled
swanlab_project: my-project  # Must be set

  5. Check Lark bot permissions: Ensure bot is added to the target group chat

Duplicate Lark notifications in multi-GPU training

Expected Behavior: Should NOT happen - only rank 0 sends notifications.

If you see duplicates: 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 RANK or LOCAL_RANK environment variables are set correctly

Solution: This is a bug if it occurs. Report with: - Full training command - Logs from all ranks - Config file

Comparison: SwanLab vs WandB

Feature SwanLab WandB
Open Source ✅ Yes ❌ No
Self-Hosting ✅ Easy ⚠️ Complex
Free Tier ✅ Generous ⚠️ Limited
Chinese Support ✅ Native ⚠️ Limited
Offline Mode ✅ Full support ✅ Supported
Integration 🆕 New ✅ Mature

Advanced Usage

Custom Logging

You can add custom metrics in your callbacks:

import swanlab

swanlab.log({
    "custom_metric": value,
    "epoch": epoch_num
})

Experiment Comparison

swanlab compare run1 run2 run3

Support

License

This integration follows the Axolotl Community License Agreement.

Acknowledgements

This integration is built on top of: - SwanLab - Experiment tracking tool - Transformers - SwanLabCallback - Axolotl - Training framework

Please see reference here

Adding a new integration

Plugins can be used to customize the behavior of the training pipeline through hooks. See axolotl.integrations.BasePlugin for the possible hooks.

To add a new integration, please follow these steps:

  1. Create a new folder in the src/axolotl/integrations directory.
  2. Add any relevant files (LICENSE, README.md, ACKNOWLEDGEMENTS.md, etc.) to the new folder.
  3. Add __init__.py and args.py files to the new folder.
  • __init__.py should import the integration and hook into the appropriate functions.
  • args.py should define the arguments for the integration.
  1. (If applicable) Add CPU tests under tests/integrations or GPU tests under tests/e2e/integrations.
Tip

See src/axolotl/integrations/cut_cross_entropy for a minimal integration example.

Warning

If you could not load your integration, please ensure you are pip installing in editable mode.

pip install -e .

and correctly spelled the integration name in the config file.

plugins:
  - axolotl.integrations.your_integration_name.YourIntegrationPlugin
Note

It is not necessary to place your integration in the integrations folder. It can be in any location, so long as it’s installed in a package in your python env.

See this repo for an example: https://github.com/axolotl-ai-cloud/diff-transformer

Adding a CLI command

An integration can contribute its own axolotl subcommand. Unlike plugin hooks, which are loaded from the plugins: list of a config file, a subcommand has to be discoverable before any config is read, so it is registered through a packaging entry point instead.

Define a regular click command in your package:

# my_package/cli.py
import click

@click.command()
@click.argument("config", type=click.Path(exists=True, path_type=str))
@click.option("--rank-ratio", type=float, default=0.01)
def my_command(config: str, rank_ratio: float):
    """one-line summary shown in `axolotl --help`"""
    from my_package.core import run  # keep heavy imports inside the function

    run(config, rank_ratio=rank_ratio)

and advertise it under the axolotl.cli_commands entry point group:

# my_package/pyproject.toml
[project.entry-points."axolotl.cli_commands"]
my-command = "my_package.cli:my_command"

Installing your package makes axolotl my-command available, and it will be listed in axolotl --help. The entry point name is the subcommand name, and the value is <module>:<attribute>.

Warning

Entry points are written into your package metadata at install time, not read from pyproject.toml at runtime. After adding or renaming one, reinstall the package (pip install -e .) or the subcommand will not appear.

Note

Commands shipped with axolotl take precedence, so a plugin cannot shadow train or any other core command. Integrations bundled in this repo are registered in BUILTIN_COMMANDS in src/axolotl/cli/plugins.py rather than through entry points, so that they work from a source checkout without a reinstall.

Adding a cloud launcher

Cloud launchers submit a complete Axolotl process before local dataset or model loading. They implement axolotl.cli.cloud.base.CloudLauncher and are discovered through the axolotl.cloud_providers entry-point group. Modal and Baseten use this same contract as built-in providers. Cloud remains an alias for CloudLauncher for existing integrations.

from axolotl.cli.cloud.base import CloudLauncher

class MyCloud(CloudLauncher):
    def train(
        self,
        config_yaml,
        launcher="accelerate",
        launcher_args=None,
        local_dirs=None,
        **kwargs,
    ):
        # Validate self.config and submit the remote Axolotl process here.
        ...

Register the provider in the plugin package’s pyproject.toml (not in Axolotl’s):

[project.entry-points."axolotl.cloud_providers"]
my-cloud = "my_package.cloud:MyCloud"

Building the plugin wheel stores this declaration in its installed .dist-info/entry_points.txt metadata. Axolotl reads that metadata at runtime; no editable checkout or pyproject.toml is required on the user’s machine. Install the plugin alongside Axolotl in the same Python environment, then select it in a separate cloud YAML:

provider: my-cloud
# Provider-specific connection, image, and resource settings go here.
axolotl train train.yaml --cloud cloud.yaml

The provider constructor receives a plain dictionary containing the cloud configuration, including provider, and owns validation of its settings. The registry calls CloudLauncher.from_config(config, config_dir=...); the default implementation calls the constructor unchanged. Override that factory when the provider needs to resolve paths relative to the cloud YAML directory. Core passes the configuration through without interpreting provider fields. Only the selected provider is imported. The built-in names modal and baseten are reserved; duplicate third-party names raise an error. Modal and Baseten are registered in Axolotl’s own wheel metadata. A built-in fallback supports source checkouts with missing or stale installed metadata. Omitting provider retains the Modal default.

For a provider that has no entry-point declaration, specify its import target directly in the cloud YAML:

provider: my_package.cloud:MyCloud

The class must be importable in the launching Python environment. This works with an existing installed package, or a local module made available through PYTHONPATH; neither requires editing or reinstalling Axolotl. The target must subclass CloudLauncher, just like a registered provider.

train receives the training YAML text, the remote process launcher (accelerate, torchrun, or python), its argument list, optional directory mounts (remote path to local path), and training configuration overrides in kwargs. Core obtains mounts by calling the provider’s get_local_dirs(cwd) hook, which defaults to no mounts. Providers own working-directory upload policy and remote mount paths. The provider must apply supported overrides and either handle or explicitly report unsupported mounts and options. Provider failures must propagate to the caller. Keep provider dependencies out of unrelated launchers.

Providers can also implement preprocess(config_yaml, *args, **kwargs) and lm_eval(config_yaml). Both operations raise NotImplementedError by default; all three cloud commands select the provider from the cloud YAML.

Provider package boundaries

Providers are bundled under axolotl.integrations but are self-contained Python packages. Each owns its configuration schemas (args.py), SDK imports, authentication, runtime templates, source staging policy, and remote helpers. Imports between files inside a provider use relative paths so the package can be moved into a standalone distribution.

Core supplies discovery, YAML reading, the CloudLauncher contract, and optional provider-independent image utilities. CloudImageConfig.from_config resolves image build paths when a provider explicitly opts into that shared schema; core dispatch does not assume providers have an image-building configuration. Training plugins and provider SDKs are not loaded for discovery.

Install provider dependencies separately while the implementations remain bundled:

pip install 'axolotl[modal]'     # Modal SDK
pip install 'axolotl[baseten]'   # Truss SDK

To extract a provider later, move its integration directory with its schemas, templates, tests, and dependency declaration into its own package. Register its launcher through that package’s axolotl.cloud_providers entry point, and remove the bundled registration/fallback for that name from Axolotl. Runtime behavior stays inside the provider; no provider-specific dispatch or schema code needs to be removed from core. The old axolotl.cli.cloud.modal_ and baseten modules currently contain only compatibility imports.

Cloud launchers and remote training backends

These are separate execution models:

Extension Submitting/client process owns Remote system owns
Whole-job launcher (Modal, Baseten, Nebius) Packaging, submission, and job monitoring The complete Axolotl training process
Training backend (Arctic Platform) Data preparation and training orchestration Model execution and optimizer state

The axolotl.cloud_providers entry-point group is scoped to whole-job launchers. Arctic-style backends use the training YAML’s plugins: list and BasePlugin hooks, including get_trainer_cls, for per-step communication with remote compute. They do not implement CloudLauncher merely because compute is remote.

The two extensions can compose: a cloud launcher can submit an Axolotl process whose trainer connects to an Arctic-style backend. Preserve the training plugin configuration when submitting a job, and provide its dependencies, connectivity, and credentials in the remote environment. The launcher API preserves that configuration; it does not define the backend’s per-step execution interface.

A launcher can also package a training plugin for remote checkpoint or completion callbacks, as needed by providers such as Nebius. Such callbacks execute inside the remote training process without initializing the training plugin manager in the submitting process.

Selecting or building the runtime image

Image preparation is separate from the launcher and training backend. Modal and Baseten accept image for a prebuilt image containing Axolotl and any custom plugins:

provider: baseten
image: registry.example.com/team/axolotl:my-fork

Alternatively, use image_build to build from an explicit local context:

provider: baseten
image_build:
  context: ../my-training-source
  dockerfile: Dockerfile
  tag: registry.example.com/team/axolotl:my-fork
  build_args:
    AXOLOTL_EXTRAS: deepspeed

context is relative to the cloud YAML file; dockerfile is relative to that context and must be inside it. Use either image or image_build. Omitting both preserves the provider’s existing default image. Builds target linux/amd64 by default; image_build.platform can select linux/arm64 for compatible providers (Modal requires linux/amd64). Build settings are validated before submission. The Dockerfile controls installation of Axolotl, dependencies, and custom plugins; building a context does not automatically install its files. For example, a context containing axolotl/ and my-plugin/ could use:

FROM axolotlai/axolotl:main-py3.11-cu128-2.9.1
COPY axolotl/ /opt/axolotl-fork/
COPY my-plugin/ /opt/my-plugin/
RUN python -m pip install --no-build-isolation /opt/axolotl-fork /opt/my-plugin

This includes local edits selected by the Dockerfile’s COPY instructions and the context’s .dockerignore. Choose a base image with a compatible Python, PyTorch, and CUDA environment, and install into the environment used by the axolotl command. Put extra package installation in the Dockerfile.

Provider image_build execution
Modal Without tag, sends the Dockerfile and context to Modal’s native image builder. With tag, builds and pushes through local Docker, then pulls the published image.
Baseten Runs local docker build and docker push, then submits the resulting registry tag. Requires Docker, a tag, and local push credentials; the remote service must also be able to pull that image.

See Modal’s Dockerfile support for supported Dockerfile instructions. Baseten’s training Image API accepts a registry image. Registry-based providers can reuse axolotl.cli.cloud.images.build_and_push_image; native builders can consume the validated CloudImageBuild settings directly. Build or push failures stop submission. This workflow builds and publishes an image when the cloud command runs; configuration parsing alone performs neither action.

Modal’s existing branch and dockerfile_commands settings remain available for its registry-image path. With image_build, put those changes in the Dockerfile; combining the modes is rejected. docker_tag remains the legacy selector for an Axolotl image tag and cannot be combined with image or image_build.

Source upload is another mechanism: Baseten’s SDK can upload a workspace and run installation commands at startup, including installing a Git fork pinned to a commit. Its training configuration does not expose a dedicated Git repository/ref field. The Axolotl Baseten launcher currently stages its config and launch scripts; use the image options above to deploy a local fork and plugins. Runtime source mounts and uploads do not by themselves install a fork or rebuild an image.

Baseten SDK compatibility

Install the tested SDK with pip install 'axolotl[baseten]' (or pip install -e '.[baseten]' from a checkout). The extra pins Truss 0.18.30. The launcher uses truss train push CONFIG, and the job definition uses the released training SDK’s Image, Compute, Runtime, and TrainingProject models. gpu selects the SDK accelerator name, for example h100 or h200.

The integration tests construct real SDK job definitions without submitting jobs. Run them in an isolated environment when updating the pin:

uv run --no-project --with truss==0.18.30 --with pytest python -m pytest \
  --confcutdir=tests/cli tests/cli/test_baseten_sdk.py

These checks cover CLI syntax, selected images, accelerator settings, secrets, and registry authentication. They do not verify live provisioning or multi-node training. Upstream documentation and the SDK’s main branch may expose fields that are not in the pinned release.

Container registries and authentication

image and image_build.tag are complete OCI image references. They can name AWS ECR, Google Artifact Registry, Azure Container Registry, GHCR, Docker Hub, or another compatible registry. Registry addresses are preserved; no Docker Hub prefix is added. Use a tag for build output; prebuilt images may use a digest.

Local builds and pushes use the Docker client’s existing login or credential helper configuration. Remote pulls use the cloud provider’s own credentials. Configure both when building into a private registry. Credentials are configured through the providers’ secret stores or workload identities, independently of the source build context. These settings concern container images; checkpoint and dataset object stores are separate provider settings.

For Modal, choose its native registry loader and a named Modal secret:

provider: modal
image: 123456789012.dkr.ecr.us-east-1.amazonaws.com/team/axolotl:fork
image_registry:
  provider: aws_ecr
  secret: ecr-pull

Use aws_ecr with an AWS IAM/OIDC secret, gcp_artifact_registry with a GCP service-account secret, or registry with REGISTRY_USERNAME and REGISTRY_PASSWORD in the named secret (for example, GHCR or ACR). See Modal private-registry authentication for secret contents. The same settings apply to a local image_build with a registry tag. Untagged native Modal builds do not use these pull credentials; for a Dockerfile with private base images, use the tagged local build path with Docker credentials for those bases.

For Baseten, docker_auth follows its SDK’s DockerAuth configuration. For example, ECR with named Baseten secrets:

provider: baseten
image: 123456789012.dkr.ecr.us-east-1.amazonaws.com/team/axolotl:fork
docker_auth:
  auth_method: AWS_IAM
  registry: 123456789012.dkr.ecr.us-east-1.amazonaws.com
  aws_iam_docker_auth:
    access_key_secret_ref:
      name: ecr-access-key
    secret_access_key_secret_ref:
      name: ecr-secret-key

This also works with image_build.tag pointing at ECR. Baseten validates the native authentication fields when Truss loads the job definition. GCP service-account JSON and generic registry secret references (for GHCR or ACR) are also supported. AWS OIDC/assume-role and GCP OIDC require a Truss release whose training DockerAuth model includes those fields; see Baseten DockerAuth. Use a Truss version that supports the chosen authentication method.