[!NOTE] LITETORCH IS COMPLETE AND PRODUCTION-READY. Officially released on PyPI: pip install litetorch.

LiteTorch

LiteTorch is a lightweight, high-performance deep learning and large language model (LLM) training engine built in native C++14 with Python bindings via pybind11.

LiteTorch delivers the core training capabilities of PyTorch + Megatron-LM + DeepSpeed ZeRO-3 within a standalone, lean C++ runtime that eliminates Python GIL latency and external dependency bloat.

Key Capabilities

1. Modern Transformer & LLM Primitives

• Rotary Position Embedding (RoPE): Native rotary embeddings as used in modern architectures like LLaMA 3 and Qwen.
• FlashAttention & GQA: Integrated FlashAttention kernel with causal masking and Grouped-Query Attention (GQA). Automatic dynamic probe for external FlashAttention-3 (libflash_attn.so) on Hopper/Blackwell.
• RMSNorm & LayerNorm: High-performance normalization layers with fused backward passes.
• SwiGLU & Activation Functions: SiLU, GELU, ReLU, LeakyReLU, Sigmoid, Tanh with warp-level GPU implementations.
• Mixture of Experts (MoE): Top-K routing with native GPU expert execution.

2. 4D Distributed Parallelism & Large-Scale Scaling

• Fully Sharded Data Parallel (FSDP / ZeRO-3): Automatic parameter, gradient, and optimizer state sharding across arbitrary cluster sizes with batched ncclGroupStart/ncclGroupEnd all-gather pipelines.
• Tensor Parallelism (TP): Megatron-LM style ColumnParallelLinear and RowParallelLinear with overlapped inter-GPU reductions.
• Pipeline Parallelism (PP): 1F1B (One-Forward-One-Backward) schedule to minimize pipeline bubbles.
• Context Parallelism (CP) & Ring Attention: Sequence splitting across GPUs with ring-based Key-Value communication.
• Rendezvous System: Dual initialization via Shared FileStore (LITETORCH_RENDEZVOUS_FILE) and TCP sockets for massive GPU scale (1000+ GPUs).

3. Dual-Platform Hardware Acceleration

• NVIDIA CUDA: cuBLAS, cuDNN, 5th-Gen Blackwell Tensor Core support (sm_100), native FP8/FP4 matrix multiplication (cublasLtMatmul).
• AMD ROCm / HIP: Full hipcc compilation with rocBLAS and MIOpen support.
• OpenCL & CPU Fallback: Automatic hardware detection falling back to OpenCL or multi-threaded CPU execution.

4. Advanced Memory Management

• Activation Checkpointing: Recomputes intermediate layer activations on the backward pass to reduce activation VRAM by 60% to 70%.
• LRU Memory Eviction & Caching Allocator: Smart block caching with automatic LRU swapping between Host RAM and GPU VRAM.
• Mixed Precision (AMP): Automatic FP16/BF16/FP8 training with dynamic loss scaling via GradScaler.
• CUDA Graph Capture: Stream recording to eliminate host-device launch latency.

Installation

Install via PyPI

pip install litetorch

Install from Source

git clone https://github.com/nguyenminh20000/Litetorch-.git
cd Litetorch-
pip install -r requirements.txt
pip install -e .

Architecture Overview

graph TD
    A["Python API (import litetorch as lt)"] --> B["C++ Binding Layer (pybind11)"]
    B --> C["Core Tensor & Autograd DAG Engine"]
    C --> D["Memory Management (LRU Eviction, Caching Allocator, Checkpointing)"]
    C --> E["Distributed Engine (FSDP, ZeRO-3, TP, PP, CP, NCCL/RCCL)"]
    D --> F["Compute Backends"]
    E --> F
    F --> G1["NVIDIA CUDA Backend (cuBLAS, cuLt, FlashAttention, Blackwell)"]
    F --> G2["AMD ROCm Backend (rocBLAS, MIOpen)"]
    F --> G3["OpenCL GPU Backend"]
    F --> G4["Multi-Threaded CPU Engine"]

Code Examples

1. Basic Tensor Operations & Autograd

import litetorch as lt

device = lt.auto_device()

x = lt.Tensor.from_vector([1.0, 2.0, 3.0, 4.0], [2, 2], device, True)
w = lt.Tensor.from_vector([0.5, -1.0, 2.0, 0.1], [2, 2], device, True)

y = lt.Ops.matmul(x, w)
loss = lt.Ops.sum(y)
loss.backward()

print("Loss:", loss.item())
print("Gradient of X:", x.grad.to_vector())

2. Transformer Decoder Layer (Self-Attention + RMSNorm + Linear)

import litetorch as lt

class TransformerDecoderBlock(lt.nn.Module):
    def __init__(self, hidden_dim, num_heads):
        super().__init__()
        self.norm1 = lt.nn.RMSNorm([hidden_dim])
        self.norm2 = lt.nn.RMSNorm([hidden_dim])
        self.q_proj = lt.nn.Linear(hidden_dim, hidden_dim, False)
        self.k_proj = lt.nn.Linear(hidden_dim, hidden_dim, False)
        self.v_proj = lt.nn.Linear(hidden_dim, hidden_dim, False)
        self.out_proj = lt.nn.Linear(hidden_dim, hidden_dim, False)
        self.fc1 = lt.nn.Linear(hidden_dim, hidden_dim * 4, False)
        self.fc2 = lt.nn.Linear(hidden_dim * 4, hidden_dim, False)
        self.hidden_dim = hidden_dim
        self.num_heads = num_heads

    def forward(self, x):
        h = self.norm1.forward(x)
        q = self.q_proj.forward(h)
        k = self.k_proj.forward(h)
        v = self.v_proj.forward(h)
        attn_out = lt.Ops.flash_attention(q, k, v, self.num_heads, self.num_heads, True)
        x = lt.Ops.add(x, self.out_proj.forward(attn_out))
        
        h2 = self.norm2.forward(x)
        mlp_out = self.fc2.forward(lt.Ops.silu(self.fc1.forward(h2)))
        out = lt.Ops.add(x, mlp_out)
        return out

    def parameters(self):
        return (
            self.norm1.parameters() + self.norm2.parameters() +
            self.q_proj.parameters() + self.k_proj.parameters() +
            self.v_proj.parameters() + self.out_proj.parameters() +
            self.fc1.parameters() + self.fc2.parameters()
        )

3. Mixed Precision Training with GradScaler & AdamW

import litetorch as lt

device = lt.auto_device()
model = TransformerDecoderBlock(128, 4)
model.to(device)

optimizer = lt.optim.AdamW(model.parameters(), lr=1e-3, weight_decay=0.01)
scaler = lt.amp.GradScaler(init_scale=65536.0)

x = lt.Tensor.from_vector([0.1] * (2 * 16 * 128), [2, 16, 128], device, False)
target = lt.Tensor.from_vector([0.0] * (2 * 16 * 128), [2, 16, 128], device, False)

for step in range(50):
    optimizer.zero_grad()
    with lt.amp.AutocastGuard(True):
        logits = model.forward(x)
        loss = lt.Ops.mse_loss(logits, target)
    
    scaler.scale(loss).backward()
    scaler.step(optimizer)
    scaler.update()
    
    if step % 10 == 0:
        print(f"Step {step} | Loss: {loss.item():.6f}")

4. Fully Sharded Data Parallel (FSDP) Training

import litetorch as lt

class LargeModel(lt.nn.Module):
    def __init__(self):
        super().__init__()
        self.layer1 = lt.nn.Linear(1024, 4096, True)
        self.layer2 = lt.nn.Linear(4096, 1024, True)

    def forward(self, x):
        h = lt.Ops.relu(self.layer1.forward(x))
        return self.layer2.forward(h)

model = LargeModel()
lt.distributed.FSDP.fully_shard(model)

Verification & Benchmarks

License

LiteTorch is released under the MIT License.