Every fine-tuning toolbox has a philosophy. axolotl bets on YAML, LLaMA-Factory on a hundred CLI flags, unsloth on monkey-patching, and ms-swift on mixin-composed trainers. LMFlow, from the Statistics and Machine Learning Research Group at HKUST, bets on something subtler: a small set of stable abstractions - model, dataset, and a family of aligner and tuner objects - with everything else pushed to the edges. The result is Apache-2.0 licensed, version 1.1.0, and notable for shipping its own research: RAFT (Reward rAnked FineTuning) and LISA (Layerwise Importance Sampling) were born here, the latter letting a 7B model train in 24GB without offloading. A major rewrite in July 2025 moved the stack onto Accelerate with optional extras for vLLM, SGLang, and TRL, and the project took the Best Demo Paper Award at NAACL 2024 along the way.

As always in this series, what follows is an educational tour of published source code.

Overview: examples and dataclass arguments on the left, the dataset and conversation-template layer, the model abstraction, the tuner-aligner-inferencer stages, and the optimizer and protocol libraries underneath.

Reading the overview from left to right:

• examples/finetune.py is the whole driver story in 60 lines: main at line 32 parses three argument dataclasses, requests a Finetuner from the package’s stage dispatcher, wraps a Dataset, loads a model through AutoModel, and calls tune at line 56.
• lmflow/args.py is the 1418-line argument universe: OptimizerNames at line 34, ModelArguments at 58, DatasetArguments at 418, FinetunerArguments at 626, InferencerArguments at 853, and AutoArguments at 1408, which maps a stage name to its argument class.
• lmflow/datasets/dataset.py defines Dataset at line 53, a backend-wrapping container with from_dict at 150, map at 309, sample at 397, train_test_split at 427, and sanity_check at 483.
• lmflow/models/auto_model.py is a 22-line dispatch: AutoModel at line 10 routes arch_type “decoder_only” to HFDecoderModel and “text_regression” to HFTextRegressionModel.
• lmflow/models/hf_model_mixin.py holds HFModelMixin at line 43, 608 lines of private preparation - __prepare_quant_config at 203, __prepare_peft_config at 248, __prepare_model_for_training at 323, activate_model_for_inference at 518.
• lmflow/models/hf_decoder_model.py defines HFDecoderModel at line 57 with tokenize at 81, inference at 313, dedicated __vllm_inference at 430 and __sglang_inference at 456 branches, and merge_lora_weights at 577.
• lmflow/optim/utils.py provides create_customized_optimizer at line 10, which synthesizes a trainer subclass whose get_optimizer_cls_and_kwargs at line 13 selects from a 22-module optimizer zoo including Muon at line 41 of muon.py and SophiaG at line 7 of sophia.py.
• lmflow/utils/protocol.py defines DataProto at line 257, a 1067-line tensor-container module with union at 688, make_iterator at 707, and chunk at 762 - the currency for handing generation results into the next alignment round.

Why You Need This

The layering is the draw. LMFlow’s BaseModel at line 5 of lmflow/models/base_model.py is six lines long, and the HF-specific implementation arrives as a mixin you compose rather than a hierarchy you inherit: HFDecoderModel at line 57 of hf_decoder_model.py combines DecoderModel, HFModelMixin, and a six-line Tunable interface from lmflow/models/interfaces/tunable.py. If you have read trl’s trainer-centric design or ms-swift’s factory-composed trainers, this is a third point in the design space - closer to a library than a framework, which is exactly why its aligners are so easy to read.

The alignment catalog is the second reason. DPO, DPOv2, iterative DPO, RAFT, and reward-model tuning each live in their own stage class, and LMFlow does not pretend one trainer fits all: the RAFT aligner carries its own vendored 3786-line Trainer copy (RaftTrainer at line 209 of raft_trainer.py) because the algorithm predates transformers 4.35, while DPOv2 rides TRL’s DPOTrainer with a custom dpo_loss at line 102 of dpov2_trainer.py and a reward-length-penalty pairing scheme. Iterative DPO loops target-model sampling, reward scoring, and preference-pair construction in a single align method. Compare this with OpenRLHF’s single-controller PPO or axolotl’s GRPO trainers, and the algorithm-as-a-class style becomes very legible - each aligner’s align() is a readable script for one research idea.

Third, the edges of the package are full of genuinely useful tools you rarely see in one place: LISA’s DynamicLayerActivationCallback at line 8 of lisa_trainer.py is 66 lines that freeze and unfreeze layer ranges on a schedule; the optimizer zoo ships 22 implementations from AdaBelief to LARS to Muon with a README benchmark table; the Evaluator at line 46 of evaluator.py implements accuracy and perplexity paths over Accelerate and DeepSpeed; and the inferencer family includes speculative decoding (SpeculativeInferencer at line 318) and a tool-calling variant whose code_exec at line 628 runs model-emitted Python. For benchmarking the trained result, the README itself points to lm-evaluation-harness, which we toured earlier, and for LoRA internals you can pair this tour with peft and unsloth.

Detailed view: the entry dispatcher with extras-conditional registration, data and template machinery, the model abstraction, tuning with LISA, the alignment algorithm family with its vendored trainer, inference backends, and the optimizer zoo.

A run begins in examples/finetune.py: main at line 32 resolves the stage-specific argument class through AutoArguments at line 1408 of lmflow/args.py, then lets HfArgumentParser fill ModelArguments at 58, DatasetArguments at 418, and the stage arguments in one pass. The dispatcher that turns the name “finetuner” into a Finetuner object is deliberately boring - a name-to-class registry - but its registration block is where the extras system lives: the SGLang and vLLM inferencers join the registry only when their packages import (the availability probes are is_vllm_available at line 60 and is_sglang_available at line 64 of lmflow/utils/versioning.py), the DPO aligners require TRL, and the RAFT aligner is registered only below transformers 4.35, matching the vendored trainer it depends on. The same pattern drives the model side: AutoModel.get_model at line 12 of auto_model.py reads arch_type and returns HFDecoderModel or HFTextRegressionModel, nothing else.

The model layer is a study in mixin composition. HFModelMixin at line 43 of lmflow/models/hf_model_mixin.py runs a chain of private preparers - __prepare_tokenizer at 94, __prepare_dtype at 133, __prepare_model_config at 159, __prepare_quant_config at 203 for the QLoRA paths, __prepare_peft_config at 248, then __prepare_model_for_training at 323 or __prepare_model_for_inference at 386 depending on the phase - and exposes activate_model_for_inference at 518 and deactivate_model_for_inference at 556 so stages can borrow the model between rounds. HFDecoderModel at line 57 of lmflow/models/hf_decoder_model.py adds the tokenizer-adjacent work: tokenize at 81, encode at 234, decode at 271, inference at 313 with __vllm_inference at 430 and __sglang_inference at 456 as dedicated backend branches, prepare_inputs_for_inference at 481, and checkpoint utilities merge_lora_weights at 577 and get_peft_without_qlora at 590. The regression sibling at line 45 of lmflow/models/hf_text_regression_model.py is the 424-line reward-model class, and the multimodal path pairs CustomAutoVision2SeqModel at line 31 of lmflow/models/vision2seq_model.py with a 294-line CLIP wrapper in lmflow/models/vision_encoder/clip_encoder.py.

Data flows through one container class. Dataset at line 53 of lmflow/datasets/dataset.py wraps a Hugging Face or JSON backend behind from_dict at 150, create_from_dict at 235, map at 309, sample at 397, and train_test_split at 427, with sanity_check at 483 catching malformed instances before they reach the trainer; the multimodal sibling in lmflow/datasets/multi_modal_dataset.py adds 291 lines for image-text corpora. Tokenization lives in lmflow/tokenization/hf_decoder_model.py - blocking at 19, tokenize_function at 57, and conversation_tokenize_function at 111, which consults the template layer. That layer, lmflow/utils/conversation_template/base.py, models a chat format as TemplateComponent lists (line 18) built by formatters (StringFormatter at 115), with ConversationTemplate at 154 providing encode_conversation at 180 and post_process_pairs at 299, and ConversationTemplateForTool at 403 adding _handle_tools at 542 for function-calling formats. Thirteen per-family modules subclass it - qwen.py at 414 lines, deepseek.py at 244, llama.py with Llama2ConversationTemplate at line 16 - and the Finetuner consumes the result through group_text at line 131 of finetuner.py, which packs tokenized sequences to model_max_length exactly the way axolotl’s packing and ms-swift’s padding-free batches do.

The alignment family is where the research lives. Finetuner at line 40 (its tune method at 222) handles full, LoRA, and QLoRA runs, and LISA plugs in as DynamicLayerActivationCallback at line 8 of lisa_trainer.py: freeze_all_layers at 43, on_step_begin at 49, and switch_active_layers at 54 rotate a budget of active layers on an interval - the layerwise importance sampling from the group’s own paper, and a cousin of the memory tricks in deepspeed and colossalai. DPOAligner at line 72 of dpo_aligner.py pairs a policy and reference through TRL, while DPOv2Aligner at line 30 of dpov2_aligner.py goes deeper: convert_to_paired_dataset at 143, _calc_reward_with_length_penalty at 198, and sampling_paired_idx_from_rewards at 210 build chosen/rejected pairs by reward, with MemorySafeDPOv2Aligner at 273 wrapping the same logic in a subprocess. IterativeDPOAligner at line 29 of iterative_dpo_aligner.py orchestrates rounds - _align_single_iteration at 66 calls _do_target_model_inference at 120, then _do_reward_model_inference at 185, then _do_single_dpo_align at 210, passing DataProto batches between steps. RaftAligner at line 34 of raft_aligner.py takes the other road: _get_batch_dataset_top at 257 ranks samples by reward-model score, keeps the top fraction, and simply fine-tunes on winners via align at 545 - no preference pairs at all - which is why it carries the vendored RaftTrainer at line 209 of raft_trainer.py, 3786 lines with train at 1517 and its own get_optimizer_cls_and_kwargs at 1059. The reward models themselves are trained and served by RewardModelTuner at line 16 of rm_tuner.py and RewardModelInferencer at line 41, both sitting on the regression model class.

Serving and evaluation close the loop. Inferencer at line 45 (with create_dataloader at 99, inference at 128, and streaming at 270) is the HF-backend default; SpeculativeInferencer at 318 implements draft-then-verify decoding with autoregressive_sampling at 422; ToolInferencer at 565 answers with tool calls and executes Python via code_exec at 628. VLLMInferencer at line 37 of vllm_inferencer.py adds batch throughput with _parse_args_to_sampling_params at 51 and a memory-safe subprocess variant at 122, returning results as DataProto at line 257 of lmflow/utils/protocol.py - a tensor-dict container whose chunk at 762, concat at 815, repeat at 869, and sample_level_repeat at 952 make it a smaller cousin of the data movement in OpenRLHF’s experience buffers. Evaluator at line 46 of evaluator.py finishes the story with evaluate at 142 dispatching to _evaluate_acc_with_accelerate at 179, _evaluate_acc_with_deepspeed at 287, _evaluate_ppl at 391, and _evaluate_nll at 440, using answer_extraction at line 155 of lmflow/utils/data_utils.py to grade free-form outputs.

From Install to a Fine-Tuned Model

The README’s recommended setup targets Linux and pins version 1.0.0 or later:

git clone -b v1.0.0 https://github.com/OptimalScale/LMFlow.git
cd LMFlow
conda create -n lmflow python=3.9 -y
conda activate lmflow
conda install mpi4py
pip install -e .

Then a LoRA run on Llama-2-7B with a conversation template:

cd data && ./download.sh alpaca && cd -
bash ./scripts/run_finetune_with_lora.sh \
  --model_name_or_path meta-llama/Llama-2-7b-hf \
  --dataset_path data/alpaca/train_conversation \
  --conversation_template llama2 \
  --output_model_path output_models/finetuned_llama2_7b_lora

Swap the script for run_finetune_with_lisa.sh and add –lisa_activated_layers 1 –lisa_interval_steps 20 to train with LISA’s rotating layer budget, or run_finetune_with_custom_optim.sh to pick any of the 22 optimizers. The Python API is the five-call sequence from examples/finetune.py - argument classes, dispatcher, Dataset, AutoModel, tune - and extras like pip install -e ".[vllm,trl]" unlock the vLLM-backed iterative DPO loop.

Honest limits: the README states the package is fully tested on Ubuntu and only partially elsewhere; the RAFT aligner is pinned to transformers below 4.35 because of its vendored trainer; vLLM and SGLang must live in separate environments due to CUDA conflicts; and WandB reporting is on by default until you pass –report_to none. The abstraction seams are visible, though - every algorithm is one aligner class you can read top to bottom, which is rarer than it should be.

Enjoyed this post? Never miss out on future posts by following us