GitHub - DLYuanGod/MegaTrain

🚀 News

• 4/17/2026: Multi-GPU data parallelism — 4x H100 achieves 4x linear speedup over single GPU. No NCCL required — workers read from shared memory independently. Thanks to @ckgresla for the initial multi-GPU PR that inspired this implementation!
• 4/12/2026: Fully integrated with the VERL framework — single H100 GPU GRPO training for Qwen3.5-27B. See RL Training.

Features

• Single GPU, Massive Models -- Train 120B+ models on one GPU by leveraging CPU RAM for parameter storage
• Multi-GPU Data Parallelism -- Scale to multiple GPUs with super-linear speedup via spawn-based workers (no NCCL)
• Universal Model Support -- Any HuggingFace decoder-only model works out of the box via AutoModelForCausalLM
• Hybrid Architecture -- Automatic handling of mixed attention (linear + full) and MoE layers
• LlamaFactory-style Data -- Flexible dataset_info.json registry with alpaca/sharegpt format support
• 1.84x Faster -- Outperforms DeepSpeed ZeRO-3 on 14B models through pipelined double-buffered execution
• YAML Configuration -- Easy model/dataset/hyperparameter setup with 25+ ready-made configs

Quick Start

# Install
git clone https://github.com/DLYuanGod/MegaTrain.git
cd MegaTrain
pip install -e .

# SFT: Train with built-in demo data
python examples/sft/train.py --config examples/sft/configs/llama3_8b.yaml

# SFT: Train any supported model
python examples/sft/train.py --config examples/sft/configs/qwen3_5_27b.yaml

# SFT: Multi-GPU data parallel (no NCCL needed)
python examples/sft/train.py --config examples/sft/configs/qwen_7b_4gpu.yaml

# RL (GRPO): Single-GPU GRPO via VERL + MegaTrain + SGLang
CUDA_VISIBLE_DEVICES=0 bash examples/rl/run_qwen2_5_7b_megatrain.sh

Supported Models

MegaTrain uses HuggingFace's AutoModelForCausalLM / AutoModelForImageTextToText with automatic model structure discovery. Both LLM and VLM models are supported without code changes. Vision encoders are CPU-offloaded just like decoder layers — GPU only holds what's currently computing.

Data Preparation

MegaTrain supports a LlamaFactory-compatible data system with flexible format support.

Option 1: Dataset Registry (Recommended)

Register datasets in data/dataset_info.json and reference by name:

dataset:
  name: "alpaca_en_demo"    # name from dataset_info.json
  dataset_dir: "data"
  max_seq_len: 1024

Supports alpaca format, sharegpt format, local JSON/JSONL files, and HuggingFace Hub datasets. See data/README.md for details.

Option 2: Direct Path (Legacy)

dataset:
  path: "/path/to/arrow/dataset"
  query_field: "query"
  response_field: "response"

Provided Datasets

Configuration

python scripts/calc_resource.py

model:
  name: "Qwen/Qwen3.5-27B"
  dtype: "bfloat16"
  attn_implementation: "flash_attention_2"

dataset:
  name: "metamath"
  max_seq_len: 1024

training:
  batch_size: 64       # <-- Use calc_resource.py to determine this!
  num_steps: 500
  learning_rate: 1.0e-5

optimizer:
  type: "deepspeed_adam"

See examples/sft/configs/ for ready-made SFT configurations. See examples/rl/ for GRPO training scripts (VERL + MegaTrain).

RL Training (GRPO)

MegaTrain supports single-GPU RL post-training via GRPO (Group Relative Policy Optimization), fully integrated with the VERL framework.

Architecture

On a single GPU, three components coexist without weight reloading:

MegaTrain stores all parameters and optimizer states in CPU RAM (~12 GB per 1B params). The GPU only holds one layer at a time during training, while SGLang's FP8 model weights stay resident for fast rollout generation.

Quick Start

# Qwen2.5-7B — recommended starting point (fast iteration, fits easily on 80GB)
CUDA_VISIBLE_DEVICES=0 bash examples/rl/run_qwen2_5_7b_megatrain.sh

# Qwen3.5-27B — full-scale single-GPU GRPO
CUDA_VISIBLE_DEVICES=0 bash examples/rl/run_qwen3_5_27b_megatrain.sh

Use a local model path instead of downloading from HuggingFace:

MODEL_PATH=/path/to/Qwen2.5-7B \
CUDA_VISIBLE_DEVICES=0 bash examples/rl/run_qwen2_5_7b_megatrain.sh

Override any parameter via Hydra CLI:

CUDA_VISIBLE_DEVICES=0 bash examples/rl/run_qwen2_5_7b_megatrain.sh \
    data.train_batch_size=16 \
    actor_rollout_ref.rollout.n=4 \
    actor_rollout_ref.actor.optim.lr=5e-7

Tested Configurations (Single H100 80GB, GSM8K)

Configuration Reference

All parameters are standard VERL Hydra configs. Key MegaTrain-specific knobs:

# Use MegaTrain as training backend (actor + reference)
model_engine=megatrain
actor_rollout_ref.actor.strategy=megatrain
actor_rollout_ref.ref.strategy=megatrain

# MegaTrain engine tuning
actor_rollout_ref.actor.megatrain.checkpoint_interval=4    # gradient checkpoint every N layers
actor_rollout_ref.actor.megatrain.num_grad_slabs=12        # async gradient buffer count
actor_rollout_ref.actor.megatrain.max_seq_len=1536         # max sequence length

# SGLang FP8 rollout
actor_rollout_ref.rollout.name=sglang
actor_rollout_ref.rollout.quantization=fp8
actor_rollout_ref.rollout.gpu_memory_utilization=0.5       # fraction for KV cache (after model weights)

See examples/rl/ for ready-to-run scripts. See verl/workers/engine/megatrain/ for the engine implementation.

Key Techniques

• Double buffering for overlapped CPU→GPU weight transfer
• Per-layer structure grouping for hybrid/MoE architectures
• Gradient checkpointing every K layers to reduce GPU memory
• Async gradient collection with slab pool and worker thread
• FP8 quantized rollout via SGLang for 2x memory savings on inference
• ref_in_actor pointer swap — zero-cost reference log-prob computation
• HuggingFace native Flash Attention integration
• DeepSpeed CPUAdam for 5-7x faster optimizer steps

Installation

git clone https://github.com/DLYuanGod/MegaTrain.git
cd MegaTrain
pip install -e .

# Optional: faster attention & optimizer
pip install flash-attn
pip install flash-linear-attention causal-conv1d  # for Qwen3.5 linear attention
pip install deepspeed                              # for CPUAdam optimizer

# Required for RL (GRPO) training
pip install verl                                   # VERL framework
pip install sglang[all]                            # SGLang rollout engine

Troubleshooting

Citation

If you use MegaTrain in your research, please cite:

@misc{yuan2026megatrainprecisiontraining100b,
      title={MegaTrain: Full Precision Training of 100B+ Parameter Large Language Models on a Single GPU}, 
      author={Zhengqing Yuan and Hanchi Sun and Lichao Sun and Yanfang Ye},
      year={2026},
      eprint={2604.05091},
      archivePrefix={arXiv},
      primaryClass={cs.CL},
      url={https://arxiv.org/abs/2604.05091}, 
}

Acknowledgement

This project benefits from the following open-source works:

• LLaMA-Factory -- Our data loading system (dataset_info.json registry, alpaca/sharegpt format support) is inspired by LlamaFactory's elegant dataset management design. Thanks to @hiyouga and all contributors.
• VERL -- RL post-training framework. MegaTrain integrates as a VERL training backend for single-GPU GRPO/PPO/DPO training.
• HuggingFace Transformers -- Universal model loading and native Flash Attention integration.
• DeepSpeed -- SIMD-accelerated CPUAdam optimizer.
• Flash Attention -- Memory-efficient attention and cross-entropy loss.
• Flash Linear Attention -- Efficient linear attention kernels for hybrid models like Qwen3.5.

License

This repository is licensed under the Apache-2.0 License.