inclusionAI/TorchTurbo

inclusionAI/TorchTurbo

Description: A supernode-native post-training engine for rapidly adapting and training LLMs & VLMs

Language: Python

License: Apache-2.0

Stars: 21

Forks: 7

Open issues: 1

Created: 2026-10-09T06:13:49Z

Pushed: 2026-10-16T06:26:06Z

Default branch: main

Fork: no

Archived: no

README:

TorchTurbo

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//
├── __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.

Vision balancing is configured on a compatible module's plan entry. For Bailing V3 VL, enable it on model.visual in the file passed to --tt_config:

engine_config:
parallel_plan:
plan:
model.visual:
vision_balance:
mode: patch_sq

patch weights whole media by patch count; patch_sq uses squared patch count. The model's Python plan supplies VisionBalanceSpec(io=BailingVisionBalanceIO()); its default mode=None leaves balancing disabled. A mode-only file override preserves that I/O instance. Python callers use VisionBalanceMode.PATCH or VisionBalanceMode.PATCH_SQ; file values are converted to the enum when parsed.

Models implement BaseVisionBalanceIO to read and replace call inputs, extract and restore all output streams, and build valid dummy media. VisionBatch declares packed patch rows, media grids, and per-medium lengths for each output stream. register_vision_balance uses the instance to preserve the model's argument and return formats while exchanging media and restoring their order. The boundary must be called in the same order on every WORLD rank. Bailing V3 VL provides this implementation; other towers require their own I/O contract.

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