LAZY-Train turns memory-limited single-GPU LLM training into an explicit resource allocation problem. It schedules layer checkpoints, CPU activation offload, recomputation, PCIe transfers, and optional NVMe activation spills under a peak GPU-memory budget, so limited H100/H800 memory is used for useful training work instead of idle slack.
Quick Start | LAZY Alpha | Resource Scheduling | Model Compatibility | Citation
LAZY stands for Limited-resource Allocation toward Zero-waste Yield Optimization.
In plain terms: for a fixed GPU memory cap, batch size, sequence length, CPU RAM, and PCIe bandwidth, LAZY-Train asks where each layer's activation should live and when it should move. Some activations stay on GPU, some are offloaded to CPU, and some are recomputed. The schedule is solved before training and then executed by the training runtime.
The current alpha focuses on Qwen3.6-27B SFT on MetaMathQA with batch size 72. It keeps a MegaTrain-compatible baseline in the same repo so throughput, TFLOPS, loss, and EM can be compared under the same data split and optimizer.
- MegaTrain is the upstream system this repository builds on. LAZY-Train keeps MegaTrain's RAM-centric, layer-streaming training architecture and adds an optimization-based resource scheduler on top.
- LLaMA-Factory informs the data and
template conventions used here. LAZY-Train keeps the LlamaFactory-style
dataset_info.jsonworkflow and Qwen-compatible SFT formatting conventions.
LAZY-Train is not a re-release of the original MegaTrain paper. Upstream MegaTrain remains the baseline and inherited runtime substrate; the LAZY contribution is the resource-allocation scheduler, compute/communication overlap modeling, 8-bit optimizer-state path, and NVMe-aware activation-tier extension.
- 7/2026: LAZY-Train alpha release: Qwen3.6-27B / MetaMathQA / batch-72 experiments with matched MegaTrain-compatible baseline, PySCIPOpt peak70 scheduling, hybrid 8-bit optimizer-state support, and NVMe/v5 activation-tier scheduling.
- 7/2026: Modeling and experiment notes are tracked under
docs/superpowers/specs/, including the MIP scheduler, hybrid 8-bit optimizer-state path, and NVMe/v5 tiered-storage schedule.
- Optimization-based scheduling -- Uses a PySCIPOpt schedule to choose GPU checkpoints, CPU activation offload, and recomputation under a peak memory cap
- Explicit PCIe modeling -- Puts activation H2D/D2H and parameter/gradient transfer into the same bandwidth-aware resource model
- NVMe third-tier scheduling -- Adds an experimental GPU/CPU/NVMe/recompute MIP path for cases where CPU activation space is capped and colder activations can spill to local NVMe SSD
- Compute-communication overlap -- The core difference from the baseline is to schedule CPU-GPU transfers into GPU compute windows, minimizing exposed communication that would otherwise leave the GPU waiting
- Alpha baseline parity -- Provides matched MegaTrain baseline and LAZY peak70 configs for Qwen3.6-27B / MetaMathQA / batch 72
- Alpha+ hybrid 8-bit fast clip -- Adds a switchable
hybrid_adamw_8bitexperiment version that keeps DeepSpeed CPUAdam as the speed path, assigns a small parameter slice to CPU 8-bit AdamW state, and enables fast foreach grad clipping - Inherited single-GPU layer streaming -- Reuses the upstream MegaTrain-compatible CPU-RAM parameter store and layer-by-layer execution path
- Model compatibility through the inherited engine -- HuggingFace decoder-only models, hybrid attention, and MoE layers follow the upstream model-discovery path; LAZY scheduling support should still be validated per model
- LlamaFactory-style Data -- Flexible
dataset_info.jsonregistry with alpaca/sharegpt format support - YAML Configuration -- Matched baseline / LAZY alpha / 8bit_optim / NVMe configs for the Qwen3.6-27B alpha experiments
# Install
git clone https://github.com/DataArcTech/LazyTrain.git
cd LazyTrain
uv venv
uv pip install -e .
# SFT: run the matched MegaTrain-compatible baseline
GPU_ID=3 ./scripts/run_qwen3_6_27b_metamathqa_b72_baseline_tmux.sh
# SFT: run LAZY-Train NVMe/v5
GPU_ID=3 ./scripts/run_qwen3_6_27b_metamathqa_b72_lazy_peak70_nvme_tmux.shThis alpha compares the MegaTrain baseline against LAZY-Train on the same Qwen3.6-27B / MetaMathQA 70-30 setup:
| Version | Run | Config | Runner |
|---|---|---|---|
| Baseline | MegaTrain-compatible fixed checkpoint baseline | examples/sft/configs/qwen3_6_27b_metamathqa_b72_baseline.yaml |
scripts/run_qwen3_6_27b_metamathqa_b72_baseline_tmux.sh |
| Alpha | LAZY-Train peak70 PySCIPOpt schedule | examples/sft/configs/qwen3_6_27b_metamathqa_b72_lazy_peak70.yaml |
scripts/run_qwen3_6_27b_metamathqa_b72_lazy_peak70_tmux.sh |
| 8bit_optim | LAZY peak70 + hybrid 8-bit + fast grad clip | examples/sft/configs/qwen3_6_27b_metamathqa_b72_lazy_peak70_hybrid8bit.yaml |
scripts/run_qwen3_6_27b_metamathqa_b72_lazy_peak70_8bit_optim_tmux.sh |
| NVMe/v5 | 8bit_optim + GPU/CPU/NVMe/recompute MIP with cuFile/GDS activation spills; default schedule is CPU15/NVMe1 with the refhybrid path locked and conservative GDS prefetch0 reloads | examples/sft/configs/qwen3_6_27b_metamathqa_b72_lazy_peak70_nvme.yaml |
scripts/run_qwen3_6_27b_metamathqa_b72_lazy_peak70_nvme_tmux.sh |
The LAZY schedule used by the alpha run is:
schedules/qwen3_6_27b_metamathqa_b72_lazy_peak70_pcie_pool.jsonThe four current comparison versions use:
- Model:
/data1/dataarc/models/Qwen3.6-27B - Dataset split:
data/metamathqa_modelscope_70_30/trainanddata/metamathqa_modelscope_70_30/test - Batch size:
72 - Sequence length:
1024 - Optimizer: DeepSpeed CPUAdam for baseline/alpha; hybrid DeepSpeed CPUAdam + small CPU 8-bit AdamW slice for
8bit_optimandNVMe - Eval: initial
0.1 * testbefore training, then sampled10000test examples every0.1epoch - W&B project:
LazyTrain
The 8bit_optim version switches are ordinary YAML fields:
optimizer:
type: "hybrid_adamw_8bit" # off: deepspeed_adam; full slow 8-bit: cpu_adamw_8bit
eight_bit_fraction: 0.02 # target fraction assigned to CPU 8-bit AdamW state
fast_grad_clip: true # off: false, use stable reference CPU grad clippingThe 8bit_optim config keeps DeepSpeed CPUAdam for almost all parameters, uses a small
8-bit optimizer-state slice, and enables fast grad clipping. In the GPU2 one-step
smoke it improved LAZY alpha from 119.18s to 98.70s per step and reduced CPU
RSS from 361.41GB to 354.33GB. The full CPU AdamW 8-bit config remains in the
repo for memory modeling and correctness smoke tests, but it is not the default
throughput path yet.
The NVMe config inherits the 8bit_optim training settings and changes the
activation storage tier. The default NVMe/v5 schedule keeps the modeled peak70
GPU cap, caps CPU activation storage at 15GB, gives NVMe activation storage
32GB, uses lazy_nvme_mode: gds, and selects 30 GPU checkpoints, 21
CPU-offload checkpoints, and 1 NVMe-offload checkpoint with zero modeled
exposed PCIe/NVMe communication. v5 also locks the already-measured fast
hybrid8bit recompute path and GPU checkpoint set, then lets PySCIPOpt solve only
the CPU-vs-NVMe tier placement. In plain language: keep the fast path and use
NVMe only as a fallback tier when CPU activation budget is tight. Runtime NVMe
reloads now default to lazy_nvme_prefetch_distance: 0. A prefetch2 probe was
fast in a short smoke but produced grad_norm=nan at step 6 in the GPU3 full
run, so the stable default is the synchronous GDS read path. The run script
uses a per-run NVMe offload directory, removes that directory on exit by default,
and runs a safe stale-directory cleanup for old per-run directories before launch.
Cleanup logs are written beside the training log as <run>.log.cleanup.
Current Qwen3.6-27B / MetaMathQA / batch-72 results are tracked as four main versions. The NVMe lines distinguish the legacy CPU8/CPU14 probes from the current CPU15/NVMe1 default:
| Version | Status | Peak GPU / CPU | Avg step | Avg TFLOPS | Last-100 tokens/s | Test eval loss |
|---|---|---|---|---|---|---|
| Baseline | Complete, 3841 steps | 60.40GB / 339.99GB |
67.451s |
176.9 |
1075.8 |
0.2854 |
| Alpha | Partial, 1044/3841 steps | 68.84GB / 362.01GB |
61.750s |
193.2 |
1197.9 |
0.1260 at step 768 |
| 8bit_optim | Complete, 3841 steps | 68.84GB / 360.23GB |
54.374s |
219.3 |
1331.8 |
0.1187 |
| NVMe/v1 CPU8 | Legacy probe observed at step 110/3841, no nonfinite/skipped steps | 69.54GB / 343.56GB |
54.28s step 2-64; 56.30s last30 |
219.6 step 2-64; 212.0 last30 |
1358.6 step 2-64; 1311.4 last30 |
initial 1.8878 |
| NVMe/v2 CPU14 | Legacy probe observed at step 10/3841, no nonfinite/skipped steps | 68.84GB / 351.59GB |
53.28s step 2-10 |
223.7 step 2-10 |
1383.9 step 2-10 |
initial 1.8878 |
| NVMe/v5 CPU15 refhybrid gds-prefetch0 | Full run active on GPU3, step 971/3841, no nonfinite/skipped optimizer steps | 68.84GB / 354.57GB |
52.819s step 2-971; 53.386s last100 |
225.65 step 2-971; 223.24 last100 |
1396.0 step 2-971; 1381.1 last100 |
0.12625 at step 768 |
Compared with the completed baseline, alpha is about +9.2% TFLOPS on its
partial run, and the completed 8bit_optim version is about +24.0% TFLOPS. The current
8bit_optim result means hybrid_adamw_8bit with eight_bit_fraction: 0.02 and
fast_grad_clip: true; it is not the slower full CPU 8-bit AdamW path.
The active NVMe/v1/v2 runs showed that adding NVMe without preserving the known
fast recompute/GPU placement can save CPU RAM but may slow backward. NVMe/v5 is
the corrected default: it keeps hybrid8bit's measured-fast block path and GPU
checkpoint set, raises the CPU activation budget enough to avoid overusing NVMe,
then solves the CPU/NVMe split under the same peak70 and communication constraints.
The v5 runtime keeps the stable synchronous GDS reload path by default. The
prefetch2 path remains useful for profiling overlap, but it is not the default
because the GPU3 full probe saw one skipped optimizer step at train step 6.
Current conclusion: NVMe/v5 is slightly faster than 8bit_optim under the same
batch-72, peak70, Qwen3.6-27B setting, but the speedup is modest. At step2-971,
NVMe/v5 reaches 225.65 TFLOPS vs 223.66 TFLOPS for the same 8bit_optim
window, about +0.89%; test eval loss is aligned at step384/768. The more
important reason to keep NVMe is capacity and flexibility: it gives the MIP a
third activation tier for larger models, larger batches, or tighter CPU memory
budgets where GPU+CPU placement would otherwise force more recompute or fail.
Detailed source files, W&B run IDs, and notes are in
docs/superpowers/specs/2026-07-12-lazy-train-three-versions-results.zh.md.
The NVMe-tier MIP details and solved schedule are in
docs/superpowers/specs/2026-07-14-lazy-train-nvme-version.zh.md.
Prepare the MetaMathQA 70-30 split if it is not already present:
cd /data2/dataarc/wuxiaojun/RL/LazyTrain
PATH=.venv/bin:$PATH .venv/bin/python3 scripts/prepare_metamathqa_modelscope_70_30.py \
--output data/metamathqa_modelscope_70_30Start the baseline in tmux:
cd /data2/dataarc/wuxiaojun/RL/LazyTrain
GPU_ID=3 ./scripts/run_qwen3_6_27b_metamathqa_b72_baseline_tmux.shStart LAZY-Train peak70 in tmux:
cd /data2/dataarc/wuxiaojun/RL/LazyTrain
GPU_ID=2 ./scripts/run_qwen3_6_27b_metamathqa_b72_lazy_peak70_tmux.shOptional overrides:
GPU_ID=2 \
WANDB_PROJECT=LazyTrain \
WANDB_API_KEY=... \
RUN_STAMP=my_run \
./scripts/run_qwen3_6_27b_metamathqa_b72_lazy_peak70_tmux.shMonitor a run:
tmux list-sessions
tmux attach -t <session>
tail -f logs/<run>.log
tail -f logs/<run>.log.cleanup
tail -f metrics/<run>.jsonl
nvidia-smiThe runner scripts create timestamped logs, metrics, and final output
directories. They also export PATH=.venv/bin:$PATH, which is required for
DeepSpeed CPUAdam to find ninja.
NVMe cleanup defaults are conservative:
NVME_OFFLOAD_CLEANUP=1 # cleanup this run's offload dir on exit
NVME_STALE_CLEANUP_HOURS=24 # pre-launch cleanup thresholdThe pre-launch stale cleanup excludes the current run directory, skips directories
with live PID-marked activation files, and only removes direct children of
nvme_offload.
Manual dry-run cleanup:
.venv/bin/python3 scripts/cleanup_lazy_nvme_offload.py \
--root /data2/dataarc/wuxiaojun/RL/LazyTrain/nvme_offload \
--min-age-hours 24Run the 8-bit optimizer smoke test:
cd /data2/dataarc/wuxiaojun/RL/LazyTrain
GPU_ID=2 \
CONFIG=examples/sft/configs/qwen3_6_27b_metamathqa_b72_lazy_peak70_8bitopt_smoke.yaml \
NUM_STEPS=1 \
./scripts/run_qwen3_6_27b_metamathqa_b72_lazy_peak70_8bitopt_tmux.shThe full 8-bit optimizer runner refuses multi-step runs unless
ALLOW_SLOW_8BIT_FULL=1 is set. The current CPU 8-bit optimizer saves CPU
optimizer-state memory and now has a fused C++ kernel, but it is still an
experimental memory-compression path rather than the default throughput path.
This guard is intentional so long 8-bit runs are not launched by accident.
Run the hybrid 8-bit smoke test:
cd /data2/dataarc/wuxiaojun/RL/LazyTrain
GPU_ID=2 \
CONFIG=examples/sft/configs/qwen3_6_27b_metamathqa_b72_lazy_peak70_hybrid8bit_smoke.yaml \
NUM_STEPS=1 \
./scripts/run_qwen3_6_27b_metamathqa_b72_lazy_peak70_hybrid8bit_tmux.shRun the full 8bit version:
cd /data2/dataarc/wuxiaojun/RL/LazyTrain
GPU_ID=2 ./scripts/run_qwen3_6_27b_metamathqa_b72_lazy_peak70_8bit_optim_tmux.shRun the NVMe version:
cd /data2/dataarc/wuxiaojun/RL/LazyTrain
GPU_ID=3 ./scripts/run_qwen3_6_27b_metamathqa_b72_lazy_peak70_nvme_tmux.shFinal model export happens after num_steps finishes. The current alpha configs
do not save a model at every eval interval.
LAZY-Train models training as a small operations-research problem:
- Decision variables: which layer boundaries become checkpoints, which
checkpoints stay on GPU, which are offloaded to CPU, which spill to NVMe, and
where PCIe fabric plus NVMe read/write transfer work is placed. The NVMe
version uses cuFile/GDS API for GPU tensor <-> NVMe raw-file spills; explicit
cpu_stagingremains available only as a fallback mode. The 8-bit optimizer variant also sets the optimizer-state memory mode used when deriving available CPU activation space. - Hard constraints: peak GPU memory cap, CPU activation capacity, NVMe activation capacity, layer order, checkpoint reachability, shared PCIe bandwidth windows, and NVMe read/write bandwidth windows.
- Objective: minimize training-step cost from recomputation and exposed communication while keeping memory below the configured caps.
The alpha schedule is generated with PySCIPOpt and stored as:
schedules/qwen3_6_27b_metamathqa_b72_lazy_peak70_pcie_pool.jsonThe experimental 8-bit optimizer schedule is:
schedules/qwen3_6_27b_metamathqa_b72_lazy_peak70_8bitopt_pcie_pool.jsonThe default NVMe/v5 tier schedule is:
schedules/qwen3_6_27b_metamathqa_b72_lazy_peak70_nvme_cpu15_refhybrid_pcie_nvme_pool.jsonThe legacy CPU14 v2 and CPU8/NVMe11 schedules are retained for comparison at
schedules/qwen3_6_27b_metamathqa_b72_lazy_peak70_nvme_cpu14_refhybrid_pcie_nvme_pool.json,
schedules/qwen3_6_27b_metamathqa_b72_lazy_peak70_nvme_cpu14_pcie_nvme_pool.json,
and schedules/qwen3_6_27b_metamathqa_b72_lazy_peak70_nvme_pcie_nvme_pool.json.
Estimate a new model before training:
.venv/bin/python3 scripts/estimate_lazy_train_plan.py \
--model /path/to/hf/model-or-config-dir \
--gpu-memory-gb 80 \
--cpu-memory-budget-gb 360 \
--h2d-bandwidth-gbps 12 \
--d2h-bandwidth-gbps 12See docs/superpowers/specs/2026-07-10-lazy-train-resource-estimator-tool.zh.md
for Qwen 27B and GPT-OSS-120B examples.
For the current peak70 alpha, the modeled cap is 70GB on an 80GB H100, leaving headroom for allocator fragmentation, kernels, and runtime buffers.
The table below describes model families supported by the inherited MegaTrain-compatible HuggingFace loading path. It is a compatibility list, not a claim that every family below has already been benchmarked with the LAZY optimizer schedule. For a new model, run the estimator first and then validate a short smoke run before launching a full experiment.
| Model Family | Model Sizes | Architecture |
|---|---|---|
| Qwen2/Qwen2.5 | 0.5B/1.5B/3B/7B/14B/32B/72B | Dense |
| Qwen3 | 0.6B/1.7B/4B/8B/14B/32B | Dense |
| Qwen3.5 | 0.8B/2B/4B/9B/27B | Hybrid (linear+full attn) |
| Qwen3.5 MoE | 35B-A3B/122B-A10B/397B-A17B | Hybrid + MoE |
| Qwen3-Next | 80B-A3B | Hybrid + MoE |
| Llama 2 | 7B/13B/70B | Dense |
| Llama 3/3.1/3.2/3.3 | 1B/3B/8B/70B | Dense |
| Llama 4 | Scout-17B-16E/Maverick | MoE |
| Mistral | 7B | Dense |
| Mixtral | 8x7B/8x22B | MoE |
| DeepSeek (LLM/Code/R1) | 7B/16B/67B | Dense |
| Phi-3/Phi-4 | 3.8B/14B | Dense |
| Gemma 2/3 | 2B/7B/9B/27B | Dense |
| GLM-4/GLM-4.5 | 9B/32B | Dense |
| InternLM 2/2.5 | 7B/20B | Dense |
| Yi 1.5 | 6B/9B/34B | Dense |
| Baichuan 2 | 7B/13B | Dense |
| GPT-OSS | 20B/120B | Dense |
| Vision-Language Models (VLM) | ||
| Qwen2-VL/Qwen2.5-VL | 2B/7B/72B | VLM (ViT + LLM) |
| Qwen3-VL | 2B/4B/8B/32B | VLM (ViT + LLM) |
| Qwen3.5-VL | 7B+ | VLM (ViT + Hybrid LLM) |
| LLaVA/LLaVA-NeXT | 7B/13B/34B | VLM |
| InternVL 2/2.5 | 2B/8B/26B/76B | VLM |
| Gemma 3 VL | 4B/12B/27B | VLM |
| GLM-4V | 9B | VLM |
| MiniCPM-V | 2B/8B | VLM |
| Llama 4 VL | Scout/Maverick | VLM + MoE |
| Any HF decoder-only model | Any size | Auto-detected |
| Any HF VLM model | Any size | Auto-detected |
LAZY-Train inherits MegaTrain's HuggingFace
AutoModelForCausalLM/AutoModelForImageTextToTextmodel discovery. The current LAZY alpha is validated on Qwen3.6-27B SFT; other LLM/VLM families should be treated as inherited compatibility until benchmarked with a LAZY schedule.
LAZY-Train keeps a LlamaFactory-compatible data system with flexible format support.
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: 1024Supports alpaca format, sharegpt format, local JSON/JSONL files, and HuggingFace Hub datasets. See data/README.md for details.
dataset:
path: "/path/to/arrow/dataset"
query_field: "query"
response_field: "response"| Dataset | Source | Format |
|---|---|---|
| alpaca_en_demo | Built-in | Alpaca |
| MetaMathQA | HuggingFace Hub | Alpaca |
| Open-Platypus | HuggingFace Hub | Alpaca |
| MathInstruct | HuggingFace Hub | Alpaca |
| CodeAlpaca-20k | HuggingFace Hub | Alpaca |
| ShareGPT4 | HuggingFace Hub | ShareGPT |
| UltraChat-200k | HuggingFace Hub | ShareGPT |
| OpenThoughts-114k | HuggingFace Hub | ShareGPT |
| OpenR1-Math-94k | HuggingFace Hub | ShareGPT |
Caution
Do NOT guess the batch_size! Use our resource calculator to find the optimal batch size for your hardware. Wrong batch size leads to OOM or wasted GPU utilization.
python scripts/calc_resource.pymodel:
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 inherited GRPO compatibility scripts
(VERL + the MegaTrain-compatible engine). They are kept for compatibility, but
the current LAZY alpha experiments focus on SFT scheduling.
| Config | Model | Architecture |
|---|---|---|
qwen_7b.yaml |
Qwen 2.5 7B | Dense |
qwen3_8b.yaml |
Qwen 3 8B | Dense |
qwen3_5_27b.yaml |
Qwen 3.5 27B | Hybrid (linear+full attn) |
qwen3_next_80b.yaml |
Qwen3-Next 80B-A3B | Hybrid + MoE |
glm4_flash.yaml |
GLM-4.7-Flash | MoE |
llama3_8b.yaml |
Llama 3.1 8B | Dense |
gpt_oss_20b.yaml |
GPT-OSS 20B | MoE |
LAZY-Train keeps the upstream MegaTrain-compatible single-GPU RL post-training path via GRPO (Group Relative Policy Optimization) and VERL. This section is retained as compatibility documentation. It is not the main contribution evaluated by the current LAZY-Train alpha.
On a single GPU, three components coexist without weight reloading:
| Component | Where | GPU Memory |
|---|---|---|
| SGLang (FP8) | GPU — rollout inference | ~3.5 GB/B params (FP8 weights + KV cache) |
| MegaTrain-compatible training engine | CPU→GPU — actor & ref training | ~4-9 GB transient (layer-by-layer streaming) |
| VERL | Orchestration — data, advantages, logging | Minimal |
The inherited engine 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 rollout generation.
# 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.shUse 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.shOverride 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-7The RL script and Hydra strategy names still contain megatrain because they
refer to the inherited training backend implementation.
| Model | Batch Size | n | GPU Memory | Time/Step | Throughput |
|---|---|---|---|---|---|
| Qwen2.5-7B | 8 | 2 | ~62 GB | ~60s | ~120 tok/s |
| Qwen3.5-27B | 2 | 2 | ~50 GB | ~230s | ~24 tok/s |
All parameters are standard VERL Hydra configs. Key MegaTrain-compatible engine knobs:
# Use the inherited MegaTrain-compatible backend (actor + reference)
model_engine=megatrain
actor_rollout_ref.actor.strategy=megatrain
actor_rollout_ref.ref.strategy=megatrain
# MegaTrain-compatible 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.
- 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
git clone https://github.com/DataArcTech/LazyTrain.git
cd LazyTrain
uv venv
uv pip install -e .
# Optional: faster attention & optimizer
uv pip install flash-attn
uv pip install flash-linear-attention causal-conv1d # for Qwen3.5 linear attention
uv pip install deepspeed # for CPUAdam optimizer
uv pip install ninja # for fused CPU 8-bit optimizer JIT builds
# Required for RL (GRPO) training
uv pip install verl # VERL framework
uv pip install sglang[all] # SGLang rollout engineOut of Memory?
- Reduce
batch_sizein config - Increase
checkpoint_interval - Reduce
max_seq_len
Slow Training?
- Use
deepspeed_adamoptimizer (5-7x faster than PyTorch AdamW) - Install Flash Attention
- Install
flash-linear-attention+causal-conv1dfor Qwen3.5 models - Increase
num_workersfor data loading - Treat
paged_adamw_8bitas an experimental LazyTrain CPU optimizer path; it saves CPU state memory and uses a fused C++ CPU kernel when available, but DeepSpeed CPUAdam remains the default speed path
New Model Not Working?
- Ensure it's a decoder-only model (not encoder-decoder like T5)
- Check
trust_remote_code: truein config if the model requires it - Try
attn_implementation: "sdpa"or"eager"if flash attention fails
If you use LAZY-Train, please cite LAZY-Train as the primary software artifact:
@misc{wu2026lazytrain,
title = {LAZY-Train: Limited-resource Allocation toward Zero-waste Yield Optimization in Large Language Model Training},
author = {Wu, Xiaojun},
year = {2026},
url = {https://github.com/DataArcTech/LazyTrain},
note = {Alpha software release}
}This project benefits from the following open-source works:
- MegaTrain -- LAZY-Train is built on top of MegaTrain's RAM-centric, layer-streaming training architecture. We keep MegaTrain as the baseline and extend it with memory/PCIe-aware resource scheduling.
- LLaMA-Factory -- Our data loading system (
dataset_info.jsonregistry, alpaca/sharegpt format support), Qwen SFT template conventions, and optimizer alias names such aspaged_adamw_8bitare inspired by LlamaFactory's practical training workflow. Thanks to @hiyouga and all contributors. - VERL -- RL post-training framework. The inherited MegaTrain engine 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.
This repository is licensed under the Apache-2.0 License.