蚂蚁 inclusionAI 开源 TorchTurbo：面向 SuperNode 的后训练训练引擎

A supernode-native post-training engine for rapidly adapting and training large language and vision models.

Overview

Post-training large models presents several challenges. Model architectures continue to evolve and introduce innovations such as mixture-of-experts (MoE), multimodality, and sparse attention. Post-training workflows, including reinforcement learning (RL), supervised fine-tuning (SFT), continual pre-training (CPT), and direct preference optimization (DPO), must flexibly adapt to these changing architectures while combining N-dimensional parallelism for high-performance training. Rapid model iteration therefore calls for a simple, general-purpose, high-performance distributed training engine that keeps the focus on model and training logic.

At the same time, AI infrastructure is entering the SuperNode era. The powerful communication capabilities of SuperNodes provide better support for training ultra-large MoE models. Against this backdrop, we introduce TorchTurbo, a distributed training engine centered on model and training logic. TorchTurbo improves model adaptation and training efficiency through deeply optimized fused operators and SuperNode affinity, taking full advantage of low inter-node communication overhead within a SuperNode.

Built on PyTorch FSDP2 and Hugging Face Transformers, TorchTurbo combines fully sharded data parallelism (FSDP), expert parallelism (EP), and Ulysses sequence parallelism (SP) to provide flexible and efficient parallel training strategies for multimodal models. Its architecture is shown below:

Core principles:

1. Model-faithful: Use clean upstream Hugging Face model implementations, replacing modules only where necessary for high performance, such as MoE, attention, and fused operations. This approach simplifies maintenance and extension.
2. Simple parallel-training transformations: Build straightforward training-parallel strategies by adding EP and SP on top of FSDP2 to achieve high training throughput.
3. Pluggable into mainstream training frameworks: Keep the training engine focused on model and training logic.
4. SuperNode-native: Optimize the training engine for SuperNode architectures.

Model

Each model directory is responsible for three tasks:

• Registering the model type
• Implementing the TurboModel or TurboVLMModel contract and model-specific hooks
• Providing the model's default strategy through get_default_parallel_plan()

Models declare their sequence parallel input contract through get_sequence_parallel_input_slicing_spec(). Before the model call, the spec's padded fields in the global batch are right-padded to a multiple of the SP size; packed sequence metadata keeps describing the valid prefix, so attention, MoE statistics, and losses ignore the tail. The backend installs slicing at the spec's boundary before activation checkpointing and sharding. The default spec slices token IDs and positions at the model entry; VLMs slice at their media-fusion boundary. Loss targets use the same shards, while media tensors and packed sequence metadata stay global.

Typical structure:

src/torchturbo/models/<model_name>/
├── __init__.py
├── model.py
├── parallel_plan.py
├── configuration_*.py       # Optional local Transformers configuration
├── modeling_*.py            # Optional local Transformers implementation
├── lf_template.py            # Optional LlamaFactory integration
├── swift_template.py         # Optional ms-swift integration
└── npu.py / gpu.py           # Optional device-specific patches

When a model needs local copies of upstream Transformers code, its directory maintains the corresponding configuration_*.py and modeling_*.py files.

ParallelPlan

A model declares its distributed execution strategy through FSDPParallelPlan:

from torch.distributed._tensor import Shard
from torchturbo.parallel import FSDPParallelPlan, FSDPPlacementSpec, FSDPPlanSpec


def get_parallel_plan() -> FSDPParallelPlan:
    return FSDPParallelPlan(
        plan={
            "model.layers.*.mlp.experts": FSDPPlanSpec(
                placement=FSDPPlacementSpec(Shard(0), "edp"),
            ),
        },
        sp_size=1,
        ep_size=1,
    )

The wrapper class returns this plan from get_default_parallel_plan(). The engine then merges EngineConfig.parallel_plan into it, including runtime sp_size and ep_size values supplied through environment variables, a configuration file, or CLI arguments.

Backend

Currently available backend:

• fsdp2: Provides weight loading, sharding, saving, and optimizer preparation based on PyTorch FSDP2 and DTensor

Future backends should be registered through src/torchturbo/engine/backend_registry.py.

Adapter

Currently registered adapters:

• llamafactory_v0: Hooks into LlamaFactory's model loading, training, and saving workflows
• msswift: Integrates with ms-swift as a plugin

TorchTurbo currently supports training only. Evaluation, prediction, and generation are not supported, with or without sequence parallelism. The LlamaFactory and ms-swift adapters reject configurations unless eval_strategy="no", do_eval=False, do_predict=False, predict_with_generate=False, and eval_on_start=False.

For ms-swift DPO, --router_aux_loss_coef sets the default for TorchTurbo's moe.aux_loss_coeff; --tt_config and explicit --tt_moe_aux_loss_coeff take precedence. The engine's routers apply the auxiliary gradients and report aux_loss / aux_loss_scaled separately from the DPO loss.

SFT samples are always isolated, including with LlamaFactory's packing=True, neat_packing=False. train_on_prompt=True enables prompt supervision within each sample; predictions never cross sample boundaries.

The src/torchturbo/adapters/rl/ package provides reusable packed-sequence and SP helpers for reinforcement learning frameworks; it is a utilities package rather than a registered adapter.

RL callers can suspend the prepared FSDP2 optimizer between training windows.

The cpu_offload training configuration keeps optimizer steps and checkpoint save/load available. Calling engine.offload_optimizer() temporarily suspends the optimizer and rejects those operations until engine.restore_optimizer(). Restoring resumes the configured training policy, including CPU offload.

optimizer.zero_grad(set_to_none=True)
engine.offload_optimizer()
# Run inference, then release its allocations before restoring training state.
engine.restore_optimizer()

Call these synchronous, idempotent APIs on every training worker at an idle training boundary. They release AdamW, AdamWSwap, Muon, and child optimizer device state and pause QK clipping statistics, while preserving parameter and scheduler bindings. Model weights and inference resources remain caller-managed. With Muon, keep FSDP parameters in place: moving them with module.to() replaces their storage and makes the next step fail, so use full CPU offload when weights must leave the device.

Muon supports full and optimizer CPU offload, and optimizer_memory_fraction=0 selects CPU Newton-Schulz and CPU exchange. These paths are experimental until they are fully validated.

To register a LlamaFactory template for a model, add an lf_template.py file to the model directory. When an adapter is activated, it scans registered models and automatically calls the corresponding register() function.

Quick Start

Build Docker

cd docker/npu/

docker build \
--platform linux/arm64 \
-f Dockerfile_aarch64 \
-t torchturbo:aarch64 \

Platform dependency scripts are available in scripts/runs/deps/, while model-specific installers are located under scripts/runs/<model>/.

Start Training

When training is launched through an entry point provided by TorchTurbo, the training framework is invoked as a dependency.

The LlamaFactory entry point, scripts/runs/run_lf.py, explicitly activates llamafactory_v0 and then calls LlamaFactory's run_exp().
The ms-swift entry point, scripts/runs/run_swift.py, parses TorchTurbo arguments and activates the msswift plugin before dispatching the selected stage.

See the Qwen3.5 MoE training guide.

Configuration

TorchTurbo-specific configuration can be provided in three ways. The precedence order, from lowest to highest, is:

Adapter defaults < TORCH_TURBO_* environment variables < --tt_config YAML file < explicit --tt_* CLI arguments

The --tt_* arguments are defined in src/torchturbo/engine/args.py.

See scripts/runs/config_example.yaml for an example YAML configuration:

torchrun --nproc_per_node=8 scripts/runs/run_lf.py \
  --tt_config scripts/runs/config_example.yaml \
  <LLaMA-Factory arguments>

The --tt_* arguments are removed before the remaining arguments are forwarded to the training framework, so they do not interfere with the original framework's argument parser.

Run Scripts

Common training scripts are organized by model under scripts/runs/. The run_lf.py and run_swift.py files are shared entry points:

scripts/runs/
├── run_lf.py                 # LlamaFactory entry point
├── run_swift.py              # ms-swift entry point
├── config_example.yaml       # Example --tt_config file
├── deps/                     # Base dependencies
│   ├── gpu.sh                # GPU dependencies
│   ├── npu.sh                # NPU dependencies
│   ├── lf_deps.sh            # LlamaFactory dependencies
│   └── swift_deps.sh         # ms-swift dependencies
└── qwen3_5_moe/
    ├── install_lf_deps.sh    # Model runtime dependency installer
    ├── install_swift_deps.sh # ms-swift dependency installer
    ├── lf_text_sft.sh        # LlamaFactory SFT training script
    └── swift_mllm_sft.sh     # ms-swift multimodal SFT script

Example:

bash scripts/runs/qwen3_5_moe/lf_text_sft.sh

Before running a script, you will typically need to check:

• model_name_or_path
• output_dir
• TorchTurbo's PYTHONPATH
• NNODES / NPUS_PER_NODE / MASTER_ADDR / MASTER_PORT
• TORCH_TURBO_SP_SIZE / TORCH_TURBO_EP_SIZE (or --tt_sp_size / --tt_ep_size)
• TORCH_TURBO_GRADIENT_CHECKPOINTING / --tt_gradient_checkpointing
• --tt_cpu_offload or backend_config.cpu_offload in the --tt_config file

Checkpoints

Training checkpoints let you resume an interrupted run. Model weights are saved in DCP format. For inference or sharing, export the model to HF format; LoRA/DoRA weights are merged into the model during export.

checkpoint-<step>/
├── model_distcp/       # Model state
├── optimizer_distcp/   # Optimizer state
├── scheduler.pt        # LR scheduler state
├── torchturbo/         # Model and PEFT configurations
└── huggingface/        # Exported HF model

--tt_hf_export controls when to export HF models. The corresponding config field is checkpoint.hf_export:

• every (default): export an HF model with every checkpoint.
• final: export an HF model only at the end of training. Intermediate checkpoints still save DCP.

LlamaFactory saves the final HF model to <output_dir>/huggingface/. ms-swift saves it inside the last checkpoint-*/huggingface/.

To resume training, keep the entire checkpoint directory, including model, optimizer, scheduler, and training framework state. Checkpoints saved with save_only_model=true contain only the HF model and cannot be used with resume_from_checkpoint. To continue training from their weights, start a new run in a new output_dir with model_name_or_path set to <checkpoint>/huggingface. DCP loading requires all state needed by the current model and ignores extra entries. DCP also saves frozen parameters, so LoRA checkpoints need storage for the full base model.

<checkpoint>/torchturbo/checkpoint.json records the source model path, the training model config, structural engine settings (backend, mixed precision, EP size, MTP, and freeze policy), and the PEFT adapter configs. When resuming from DCP state, the current PEFT config must match the saved one, so adapter structure and LoRA scaling (lora_alpha, use_rslora, alpha_pattern) stay the same. Only dropout, bias, initialization settings, and bookkeeping fields such as peft_version and base_model_name_or_path may differ; the freeze policy may change as well. scripts/checkpoint/dcp_to_hf.py uses the file to rebuild the model. Checkpoints without this file skip the resume check and need explicit export arguments: --model_name_or_path, the --tt_* settings used for training, and for LoRA, --adapter_config or all of the script's LoRA options.

Use this script to convert a DCP checkpoint to an HF model:

torchrun --nproc_per_node=8 scripts/checkpoint/dcp_to_hf.py \
  --checkpoint_dir output/my_run/checkpoint-1000 \
  --output_dir exports/my_run/step-1000

The exported model is saved in exports/my_run/step-1000/huggingface/. The source model must remain accessible; if it has moved, pass its new path with --model_name_or_path. Export uses the saved expert parallelism (EP) size and MTP settings. The total number of processes must be a multiple of the EP size.

Use a separate export directory to avoid affecting automatic cleanup of training checkpoints. The script rejects identical input and output directories, paths where one directory contains the other, output directories that hold training checkpoint state, and output directories inside the source model directory.

Supported Models

Built-in models are discovered from src/torchturbo/models/. Public environment setup and run instructions are currently available under docs/runs/:

docs/runs/
└── qwen3_5_moe/
└── bailing_v3/
...

Environment Variables

Command-Line Arguments

Most environment variables have corresponding higher-priority --tt_* CLI arguments. All arguments are defined in src/torchturbo/engine/args.py and grouped below by category:

DeepSeek-V4 MTP and DSpark jointly train the backbone and draft modules, including the shared embedding and output head. DSpark stops gradients through its reference distributions while allowing gradients through the backbone context features.

Adding a New Model

1. Create a model directory under src/torchturbo/models/.
2. In model.py, inherit from TurboModel for language models or TurboVLMModel for multimodal models, and register the class with @register_model().
3. Define the default parallel strategy in parallel_plan.py and return it from the class method get_default_parallel_plan().
4. Add configuration_*.py and modeling_*.py if needed.
5. Add lf_template.py or swift_template.py when the corresponding framework requires model-specific registration.
6. Override hooks such as on_parallelized() or on_materialized() only when the model needs setup at those lifecycle points.
7. Add launch scripts under scripts/runs/<model>/ and a run guide under docs/runs/<model>/.

Minimal wrapper class:

from transformers import AutoModelForCausalLM

from torchturbo.parallel import FSDPParallelPlan

from .configuration_my_model import MyConfig
from .modeling_my_model import MyModelForCausalLM
from ..base import TurboModel
from ..registry import register_model


@register_model("my_model", config_cls=MyConfig, auto_class=AutoModelForCausalLM)
class TurboMyModelForCausalLM(TurboModel, MyModelForCausalLM):
    config_class = MyConfig

    @classmethod
    def get_default_parallel_plan(cls) -> FSDPParallelPlan:
        from .parallel_plan import get_parallel_plan

        return get_parallel_plan()

    def on_materialized(self) -> None:
        """Initialize runtime buffers after parameters are materialized."""

on_materialized() is optional; omit it when the model has no runtime buffers to rebuild. It runs after checkpoint loading or device materialization. Use on_parallelized() instead for setup that must happen immediately after module replacement and sharding.

Citation

If you find this project useful, please consider citing it as follows:

@misc{torchturbo2026,
  title        = {TorchTurbo: A supernode-native post-training engine for rapidly adapting and training large language and vision models.},
  author       = {Tai An and Haishan Feng and Xuejin He and Chen Long and Yicheng Shan and Wei Tang and Chao Wang and Xiong Xu and Yifei Yang and Liangji Zhuang},
  year         = {2026},
  url          = {},
  license      = {Apache-2.0}
}

This project benefits from Megatron-LM. We thank its authors for their contributions.

License

Apache-2.0