Axolotl v0.21.0: Single-node Beam tuning, expanded diffusion and multimodal training

1. Tuning a massive model on an 8-GPU server

Reflection's Beam is a mixture-of-experts model with 501 billion total parameters and 23 billion activated per token. Only some experts participate in computation for each input. It targets software development and long-running, multistep agent tasks.

This release supports QLoRA fine-tuning on a single node with 8 H200 GPUs. This method stores the original weights at low precision and trains small adjustment weights. It does not mean retraining the entire model or running it on low-end hardware.

Configurations offer FSDP2, which distributes training states, and device_map, which splits model layers across devices. The distribution method and tuning targets must be treated as separate choices.

Attention is the computation that determines which parts of the input reference one another. Both distribution methods offer configurations for tuning attention and expert weights. device_map also offers an attention-only configuration.

The device_map configuration runs in a single process. The adamw_torch_4bit optimizer, which updates weights, fails when adjustment weights span multiple devices. Use adamw_torch as specified in the documentation.

The router that selects paths must be excluded from tuning. There is a warning that training it destabilizes expert selection.

2. Diffusion models join the standard training path

Diffusion language models learn to restore masked or noisy tokens. With the new top-level diffusion_lm setting and existing conversation datasets, model-specific profiles determine the noise process, attention layout, and loss.

According to the diffusion model guide, supported models include the 3 billion- and 8 billion-parameter Nemotron-Labs-Diffusion models and Dream v0.

DiffusionGemma is also supported, but standard training for these models is limited to LoRA, which trains small adjustment weights.

Nemotron supports variable-length attention, which gathers only actual tokens from documents of different lengths for computation. This path requires Torch 2.14 or later. The documentation's installation section provides this environment setup example for that path.

python -m venv .venv-diffusion
source .venv-diffusion/bin/activate
python -m pip install --upgrade pip
python -m pip install torch==2.14.0  # 또는 가속기에 맞는 휠
python -m pip install -e .

The 3 billion-parameter starter example performs two optimizer updates using 32 synthetic conversations. It checks data preparation, training, and saving, rather than evaluating quality.

Packing multiple documents into one row primarily saves memory. The documentation recommends measuring throughput with packing both enabled and disabled, since it depends on document length and the attention implementation.

3. Applying preferences and rewards to image responses

According to the release notes, DPO and KTO can now run on image datasets. This provides a training path for aligning responses with preferences even when inputs include images.

Reinforcement learning that uses rewards for generated responses, such as GRPO, is also supported. However, no figures for image-training quality improvements were provided.

Using a generation server for image reinforcement learning also requires vLLM 0.30.0 or later.

Image input for diffusion models follows a separate path. Nemotron-Labs-Diffusion-VLM-8B accepts images only through a dedicated plugin that trains it to fill designated answer slots with labels, scores, and similar outputs. Regular supervised fine-tuning accepts only text and rejects multimodal datasets.

4. Integration benefits, runtime checks required

Standard diffusion training uses a top-level setting and familiar conversation datasets instead of specifying a plugin, reducing separate plugin configuration work. Existing diffusion configurations are converted with warnings. During migration, checking those warnings and each model's support scope remains necessary.

The B200 performance guide is a separate comparison. In small internal tests, the path using default inter-device communication without extra installation was about 25% slower than DeepEP.

Implementations of attention also differ. The same guide describes FlashAttention 2 as an alternative about 2 times slower than FlashAttention 4 at this model scale. Both figures compare implementation choices on that hardware, rather than estimating the speed or cost of the server configuration described earlier.

This release's value lies in expanding tuning for large models and training for diffusion and image tasks, while reducing separate plugin configuration work for standard diffusion training.