Skip to content

v0.18.0

Latest
Latest

Choose a tag to compare

View all tags
@NanoCode012 NanoCode012 released this 17 Jul 12:28
v0.18.0
2f5cb9d
This commit was created on GitHub.com and signed with GitHub’s verified signature.
GPG key ID: B5690EEEBB952194
Verified
Learn about vigilant mode.

Axolotl v0.18.0 Release Notes

We've been hard at work doing low level improvements in the kernels. Over 90 commits since v0.17.0 (June 3, 2026), themed around fine-tuning very large sparse-MoE models cheaply: 4-bit expert LoRA/QLoRA (NVFP4, MXFP4, bnb) that runs fast and stays memory-flat at long context on Blackwell and Hopper, plus a supply-chain hardening pass on CI.

Highlights

NVFP4 MoE-LoRA: ScatterMoE and SonicMoE

4-bit NVFP4 expert LoRA now runs fast on both MoE kernel backends.

ScatterMoE

ScatterMoE gains fused NVFP4 (Marlin/DeepGEMM) and bnb-4bit expert paths for Gemma 4 (128 experts) and DeepSeek-V4. Active VRAM stays nearly flat from 4k to 32k context: NVFP4 is fastest at short sequences, bnb uses the least memory.

scattermoe_nvfp4_gemma4

Configs: examples/gemma4/26b-a4b-moe-nvfp4-lora.yaml (ScatterMoE)

SonicMoE

SonicMoE adds native FP4-activation MoE-LoRA (W4A4: 4-bit weights and activations) through the quack/CUTLASS kernels. On B200 it beats the Marlin W4A16 (16-bit-activation) path at every sequence length at matched quality, validated on the Qwen3-30B-A3B and Qwen3-Next-80B-A3B NVFP4 checkpoints.

It also runs on consumer Blackwell (RTX 50xx / sm120) via the quack 0.6 migration (#3830), and the new nvfp4_merge_aware mode keeps the adapter bitwise-consistent when merged back into the NVFP4 base (#3822).

sweep_fp4_scale sweep_fp4_stageD-2

Configs: examples/qwen3/30b-a3b-nvfp4-lora.yaml (SonicMoE).

GLM-5.2 (DSA) Fine-tuning with 2D Expert Parallelism

Fine-tune GLM-5.2 (glm_moe_dsa) from its NVFP4 checkpoint on multi-GPU FSDP2. It follows the DeepSeek-V3.2 sparse-MLA lineage: 256 routed experts with Lightning-Indexer token selection. DeepEP expert parallelism now composes in 2D (EP × cp), and dedicated DSA attention kernels (use_glm_dsa_kernels) with rank0 + broadcast loading keep the ~250 GB 4-bit experts from OOMing on load. Config: examples/glm_moe_dsa/glm-5.2-nvfp4-lora.yaml (requires deep_ep).

Hidden-States Activation Offloading for Long-Context Full-Parameter Training

New activation_offloading: hidden_states offloads only the per-layer checkpoint input to CPU and recomputes the rest, so PCIe stays within budget where the offload-everything mode saturates it. On Qwen3-8B full-param it reaches 128k context where plain gradient checkpointing OOMs, and the memory advantage widens with sequence length (1.23x less at 64k, tied throughput) with bit-exact gradients.

Multi-Turn Inference Chat Interface

axolotl inference config.yaml --chat starts an interactive multi-turn chat with streaming token output, runtime-adjustable generation parameters, and streaming reasoning blocks for thinking models.

inference-chat

Performance & Kernel Optimizations

  • Blackwell (sm120) MoE-LoRA (#3714 by @winglian): makes the vendored ScatterMoE Triton path the working MoE+LoRA story on sm120 where the SonicMoE CUTLASS kernel can't compile. EP sentinel-skip drops masked remote rows (2 to 10x fwd+bwd at ep 2 to 8), gpt_oss layout support (9.5 to 48x vs eager), fused MXFP4 with no dequant (1.7 to 3.7x, 16 to 24x less transient memory), and a SonicMoE-to-ScatterMoE fallback on sm120.
  • Grouped-Gram dA/dB for large-E MoEs (#3712 by @winglian): recompute-free grouped-Gram LoRA weight grads plus a sync-free dX_lora path, up to 2.2x fwd+bwd on Qwen3-MoE / DeepSeek (E>=128), bit-identical to the split kernel.
  • LoRA kernel memory (#3704 by @winglian): removes an intermediate materialization in the LoRA kernel op.
  • Faster multimodal assistant-only masking (#3672 by @thad0ctor): vectorized role-boundary scanner (byte-identical to the reference) plus a fused process_labels, ~1.3 to 1.5x on Gemma 3/4 and Qwen 2 under DataLoader-worker conditions.
  • torch.compile coverage for 4-bit dequant (#3677 by @thad0ctor): registers the NF4 dequant fast path as a Dynamo-opaque torch.library.custom_op, removing the ctypes trace failure and get_ptr recompile thrash on torch_compile: true QLoRA. Eager path is byte-identical.
  • Custom torch ops for in-repo kernels (#3788 by @winglian): registers ~25 axolotl:: custom ops (attention, SwiGLU/GeGLU, RMSNorm, dsv4/glm_dsa) so kernels compile without graph breaks; a follow-up (#3789) hoists multidoc routing out of the d512 attention path to clear the last graph breaks under torch.compile.
  • Fused LoRA for GatedDeltaNet (#3732 by @thad0ctor): route Qwen3.5 GatedDeltaNet linear-attention LoRA projections through the fused kernel under lora_qkv_kernel / lora_o_kernel, avoiding a bf16/fp32 activation round-trip.

New Features

  • Multi-adapter MoE LoRA (#3719 by @winglian): multi-LoRA (multi-tenant) support for the ScatterMoE and SonicMoE kernels with optimized routing and gradient computation.
  • Selective activation checkpointing (#3786 by @winglian): selective_checkpointing saves attention outputs (SDPA/flash-attn) instead of recomputing them, fully eager, with optional CPU offload of the saved tensors.
  • Per-module LoRA rank/alpha (#3673 by @thad0ctor): lora_rank_pattern / lora_alpha_pattern set per-module LoRA rank and alpha, forwarded to PEFT and honored by the memory-efficient merger.
  • optimizer: sinkgd (#3763, #3785 by @winglian): stateless gradient multi-normalization optimizer (Scetbon et al., 2025), ~87% less optimizer-state memory on an 8B full finetune (~24% lower peak vs 8-bit AdamW) at matched loss, extended with width transfer, spectral norm, and fused Triton kernels (#3785).
  • torch_compile_options (#3692 by @thad0ctor): expose an allowlist of torch._inductor.config flags (e.g. max_autotune_gemm) via config, validated at preprocess.
  • Plugin CLI subcommands (#3840 by @winglian): integrations register axolotl subcommands via the axolotl.cli_commands entry-point group, resolved lazily so --help no longer imports torch.
  • FSDP min_num_params for size-based wrap (#3710 by @thad0ctor): expose fsdp_config.min_num_params in YAML for SIZE_BASED_WRAP.
  • Lower-RAM SFT dataset loading (#3711 by @ved1beta): fixes host-RAM OOM when loading large preprocessed SFT datasets.

Documentation

  • Support matrix (#3782 by @NanoCode012): a model/feature support matrix in the docs (closes #3440).
  • Contributor & agent testing guidance (#3748 by @thad0ctor): improved test docs.
  • DeepEP pin for expert parallelism (#3778 by @trevorgordon981): pin DeepEP to v1.2.1 in the DISABLE_NVSHMEM build instructions, since DeepEP main removed the code path the patches edit.

Model & Framework Support

Deprecations

  • Non-uv Docker images removed (#3740 by @NanoCode012): the pip-based image variants are dropped; -uv images are the supported path, with general Docker cleanup alongside.

New Model Support

  • GLM-5.2 (glm_moe_dsa) (#3759 by @winglian): see Highlights.
  • Gemma 4 Unified (encoder-free multimodal) (#3706 by @NanoCode012): text + vision LoRA configs, a dedicated chat template, and Liger + Cut Cross Entropy support.
  • DeepSeek-V4 MoE-LoRA kernels (#3747 by @winglian): fused NVFP4/bnb expert paths for DeepSeek-V4 alongside gemma4 (see Highlights).
  • MiniMax M2 (MoE) (#3702 by @ved1beta): fine-tuning support with QLoRA, kernel-optimization patches, and an example config.
  • PaddleOCR-VL (#3774 by @thad0ctor): multimodal image-text support with a dedicated processing strategy (no packing/Liger/CCE), plus QLoRA and full fine-tune examples.

Dependency Updates

  • transformers 5.9.0 → 5.14.1 (#3706, #3759, #3761, #3822): now pinned ==5.14.1.
  • trl 1.5.1 → 1.8.0 (#3784 by @ved1beta, #3822 by @NanoCode012).
  • liger-kernel 0.7.0 → 0.8.0 (#3713 by @thad0ctor): adopts liger's native auto-patch for gemma4_text / qwen3_5_text / ministral / nemotron while keeping axolotl's gated-RMSNorm handling for qwen3_5 / qwen3_5_moe; a follow-up (#3826) drops the now-shadowed llama4 / qwen3 / qwen3_moe hand-patches.
  • huggingface_hub 1.1.7 → 1.17.0 (#3706, #3804 by @NanoCode012): now pinned ==1.17.0.
  • httpx added to the uv Docker image (#3770, #3775 by @NanoCode012): pinned <1 to avoid pulling a pre-release.
  • numpy loosened to >=2.1,<3.0 (#3843 by @winglian): numba metadata governs the effective ceiling.
  • pydantic pinned ==2.12.5 (#3761, #3804).
  • cu130 builds gated to torch 2.11.0+ (#3761 by @winglian); uv Docker images for torch 2.11, 2.12, and 2.13 (#3715, #3816 by @winglian).

Bug Fixes

  • LoRA merge on quantized bases (#3771 by @winglian): dequantize quantized (NVFP4/bnb) expert bases before folding the LoRA delta so fused expert weights are preserved.
  • Gemma 4 shared-KV key on transformers >=5.8 (#3701 by @thad0ctor): key the shared K/V by layer_type instead of the removed attribute.
  • CCE update and post-merge bugs (#3757 by @NanoCode012).
  • DPO user_defined strategy KeyError on custom field names (#3742 by @vineethsaivs).
  • KTO user_defined dataset transform crash (#3730 by @Ayushhgit): fixed a crash that hit every documented config.
  • _get_messages NameError when JSON decodes to a non-list (#3739 by @JSap0914), IndexError in metharme/reflection _tokenize on empty fields (#3751 by @vineethsaivs).
  • qwen3_5 template fixes (#3728 by @ved1beta, #3725 by @Anai-Guo): fix chat handling and extract reasoning_content before reassigning content.
  • Outdated multimodal jinja template (#3736 by @NanoCode012).
  • gc_collect_steps not honored (#3709 by @thad0ctor), train_per_sec_per_gpu grad-accum miscalculation (#3699 by @NanoCode012).
  • curriculum_sampling honored in the GRPO sampler (#3707 by @lollinng, fixes #2376).
  • fsdp_version no longer flagged as a validation warning (#3718 by @SamuelLarkin).
  • numpy version mismatch (#3662 by @ved1beta).
  • Fused LoRA kernel correctness (#3799 by @winglian): fix fp32 NF4 quant-state dequantizing to garbage across the fused kernels, plus GatedDeltaNet backward dtype and partial-mixer fallback hardening.
  • Traceable fused LoRA backward (#3800 by @thad0ctor): stop LoRA MLP/QKV/QK backward writing into the saved input, restoring torch.compile fullgraph tracing.
  • Fused LoRA kernel edge cases (#3805 by @roycho96, #3806 by @thad0ctor): upcast the ScatterMoE gates-branch grad_out under bf16 autocast, and run the fused QKV/O rewrite when lora_dropout > 0.
  • FSDP2 quantized LoRA checkpoint saving (#3777, #3779 by @winglian): route quantized (NVFP4/bnb) LoRA checkpoints through the EP LoRA saver and adapter gather so sharded saves preserve the experts.
  • Sample-packing cross-document isolation (#3796 by @winglian): fix sdpa/eager packing leaking attention across documents and make sdpa_varlen engage during training.
  • Streaming pretrain long documents (#3717 by @ved1beta): chunk documents longer than sequence_len instead of dropping them.
  • Preprocess honors trust_remote_code (#3819 by @NanoCode012): respect the configured setting instead of forcing it on, and handle pre-download failures gracefully.
  • fp8: exclude MoE router linears from torchao conversion (#3791 by @ved1beta).
  • Gemma 4 GQA guard on transformers 5.13+ (#3793 by @winglian): handle the 3-arg use_gqa_in_sdpa signature.
  • Gemma end-of-turn tokens in examples (#3756 by @thad0ctor): set eot_tokens so the real turn terminator is trained.
  • KTO chatml.argilla_chat (#3838 by @vineethsaivs): read the user turn from the completion column, fixing a KeyError.
  • Diffusion padding mask (#3735 by @Ayushhgit): exclude padding from the bidirectional attention mask in non-packed diffusion training.
  • vllm-serve CLI boolean overrides (#3765 by @Anai-Guo): explicit CLI False now overrides config-enabled enable_prefix_caching / enable_reasoning.
  • Clearer jinja chat-template validation error (#3802 by @lxcxjxhx).
  • Activation-offload correctness (#3833, #3835 by @winglian): clone offset/non-contiguous saved tensors on the compute stream to fix a race corrupting shift_labels, and support the torch 2.13 checkpoint API in SAC offload.

Infrastructure

  • Supply-chain hardening: pin all GitHub Actions to commit SHAs (#3760 by @XananasX7): 78 pins across 11 workflows, replacing mutable @vX tags with immutable SHAs to close the tj-actions-style attack class. Follow-up (#3804 by @NanoCode012) sets persist-credentials: false on checkouts and adds Dependabot to keep the hashes current.
  • Fail CI early on CUDA errors (#3737 by @NanoCode012), cut the CPU test tail (#3705 by @winglian): drop dataset_num_proc to 1 and split builder tests.
  • Dependabot GitHub Actions bumps (#3810, #3812, #3813, #3814, #3821): actions/checkout 4→7, codecov-action 5→7, setup-uv 7→8, setup-buildx 3→4, actions-netlify 3→4; plus a minimum cooldown (#3820 by @NanoCode012).
  • Faster Docker builds (#3790, #3792, #3794 by @winglian): move common package and apt installs into the base image; add a torch 2.13 base image (#3816) and preinstall kernels-data before the editable build (#3841).
  • Prune the pytest matrix (#3834 by @winglian): gate expensive GPU tests behind maintainer labels, and mark a flaky VRAM-leak regression test as non-strict xfail (#3807).

New Contributors

Full Changelog: v0.17.0...v0.18.0

Contributors

winglian, thad0ctor, and 12 other contributors
Assets 2
Loading