Lazy loaded image8.3 OpenRLHF: SFT Qwen3-0.6B and Code Analysis

8.3.1 Introduction and Installation of OpenRLHF

OpenRLHF is a high-performance RLHF framework based on Ray, DeepSpeed, and HF Transformers:
  1. Easy to Use: OpenRLHF is one of the simplest and high-performance open-source RLHF libraries, fully compatible with Huggingface models and datasets.
  1. High Performance: In RLHF training, 80% of the time is spent on sample generation. Thanks to the acceleration capabilities of Ray, Packing Samples, and vLLM, OpenRLHF achieves performance gains of over 3–4x compared to DeepSpeedChat with Hybrid Engine.
  1. Distributed RLHF: OpenRLHF uses Ray to distribute the Actor, Reward, Reference, and Critic models across different GPUs while placing the Adam optimizer on the CPU. This allows efficient fine-tuning of models over 70B parameters on multiple A100 80G GPUs and vLLM, or fine-tuning 7B models on several 24GB RTX 4090 GPUs.
Additionally, OpenRLHF supports full-parameter SFT, efficient QLoRA fine-tuning, DPO algorithms, PRM training, etc.
Environment Setup:
You can create a Conda environment and install OpenRLHF by following the official quick-start guide. Colab is recommended, and the full code is available at the end.
<ins/>

8.3.2 SFT Training

We will perform supervised fine-tuning (SFT) on Qwen3-0.6B using OpenRLHF. Let’s start with an example provided by OpenRLHF, and modify it step by step for our experiment:

Prepare Dataset

Same as Chapter 8.2, we use alpaca dataset to fine-tune the base model. To lanuch the training script openrlhf.cli.train_sft, it requires to input input_key and output_key for our dataset. However, alpaca has the following format: instruction + [optional input] → output. We need to manually combine instruction and optional input as input. Then save it into disk.

Modify the OpenRLHF's Source Code for Model Saving

By default, OpenRLHF calls DeepSpeed’s save_modelmethod after training, which unconditionally uses DeepSpeed’s _consolidated_16bit_state_dict() function. However, when we run with zero_stage = 0 (because parallelization isn’t needed for a single GPU), DeepSpeed wraps the model but performs no ZeRO partitioning. The _consolidated_16bit_state_dict()function works only when weights are partitioned. Since at stage 0 there’s no partitioning, calling this function triggers an error (This is a bug, and may be fixed in the future version). To avoid this, we disable OpenRLHF’s default model-saving call and instead save the model directly using PyTorch.
  • Let’s check OpenRLHF’s SFT training script.
      • The output show that after training, OpenRLHF save model by calling strategy.save_model function.
  • We need to remove this saving method, and use torch function to save the model and tokenizer.
    • Then, we observe that the file is been correctly modified

      Training the Model By OpenRLHF

      • By default, you run deepspeed directly from the terminal. In this example, we use Python’s subprocess module to call deepspeed from within a Python script.
      • For this example, we train for only 256 iterations with a batch size of 2 to avoid CUDA out-of-memory errors.
      ⭐
      Optional: Setup wandb to Track Training Process
      • Weights & Biases (W&B) is a machine learning experiment tracking and collaboration platform. After signing up, you’ll receive a 40-digit API key.
      • The wandb package is installed automatically when you install OpenRLHF, so you only need to prepare your API key, and no additional installation steps required.
      • To enable, please uncomment --use_wandb(and provide your API key), --wandb_project, and --wandb_run_name.
      The training parameters are set as follows:
      Why we choose these settings:
      1. input_key and output_key: These match the field names defined in the dataset.
      1. Flash Attention (flash_attn) is disabled: Flash Attention only works on Ampere GPU architectures (RTX 30 series and newer). Since Google T4 GPUs do not support it, we disable it.
      1. Learning rate scheduler (lr_scheduler): By default, it uses cosine_with_min_lr, which gradually decreases the learning rate for better convergence. However, for this small-scale example, we set it to constant so the model can learn more aggressively over the short run.
      1. input_template: This must match the format used during inference. If we train with the apply_chat_template option but infer with a different format, the model may fail to recognize the correct end-of-sequence (EOS) stop token <|im_end|>.
      After training, we can open wandb to view the training process:
      notion image

      Load and Test the Model

      Since we saved the fine-tuned model and tokenizer fully to local disk, we can load them directly without downloading from Hugging Face:
      Let’s test the model with a basic prompt:
      After fine-tuning with OpenRLHF, the model correctly stops generating when it finishes the answer. Remember that in Chapter 8.2, the base model (model before SFT) doesn’t know how to generate EOS token.
      Important: The chat template at inference must match the format used during training. Since we didn’t introduce <think> and </think> tokens in training (unlike the default Qwen3-0.6B chat template), we also exclude them here during inference.
      <ins/>

      8.3.3 OpenRLHF Code Analysis

      Note: All of the following analysis is based on release version v0.7.4. Since the code may change in future releases, every link below explicitly targets v0.7.4 to ensure you can reproduce these results.

      Preparing the Dataset

      The __getitem__ method in SFTDataset (view source) is called each time you iterate over the dataset. It performs these steps:
      1. Append the <EOS> token to the end of the concatenated prompt and response (unless in pretrain mode).
      1. Tokenize the resulting text.
      1. Generate a loss mask.
      1. Return the input token IDs, attention mask, and loss mask.
      Computing the Loss Mask: In causal language model training, we shift labels one token to the left:
      • At input position , the model predicts token .
      • There is no “next” token for the final input position, so we do not provide a label (and thus no loss) for that logit.
      Collation and Padding: When batching, the collate_fn pads all sequences to the length of the longest sequence in the batch:
      📌
      Example of Data Generation
      Consider a toy vocabulary/tokenizer scenario (ignoring blank tokens):
      • Prompt: <|im_start|> How are you <|im_end|> <|im_start|>assistant
      • Model Output: I 'm doing well . <|im_end|>
      1. Prompt IDs (length 7):
          2. <|im_start|>, How, are, you, <|im_end|>, <|im_start|>, assistant
      1. Response IDs (length 6):
          2. I, 'm, doing, well, ., <|im_end|>
      1. Concatenated Input IDs (length 13). Suppose the batch’s max length is 15, two <|pad_token|> tokens are appended.
      Token
      ID
      Attn Mask
      Model’s Predicts
      Loss Mask
      <|im_start|>
      1
      1
      How
      0
      How
      101
      1
      are
      0
      are
      102
      1
      you
      0
      you
      103
      1
      <|im_end|>
      0
      <|im_end|>
      2
      1
      <|im_start|>
      0
      <|im_start|>
      1
      1
      assistant
      0
      assistant
      104
      1
      I
      1
      I
      105
      1
      'm
      1
      'm
      106
      1
      doing
      1
      doing
      107
      1
      well
      1
      well
      108
      1
      .
      1
      .
      109
      1
      <|im_end|>
      1
      <|im_end|>
      2
      1
      ???
      0
      <|pad_token|>
      3
      0
      ???
      0
      <|pad_token|>
      3
      0
      ???
      0
      • Attention mask is 1 for all non-padding tokens.
      • Loss mask is 1 for all response tokens (i.e., those the model should learn to predict).

      Loss Calculation in SFT Training

      The key training code starts at line 135 of the sft_trainer.py file. A simple version (after removing auxiliary features) is the following:
      The SFT Loss Function
      Here, masked_mean multiplies the negative log‐probs by the loss mask (1 for tokens to include, 0 otherwise), sums them, and divides by the total count of masked tokens:
      Putting it all together, if we index the batch by and tokens by , the SFT loss is
      where:
      • is the batch size,
      • is the number of prompt tokens in example
      • is the number of response tokens,
      • is the model’s predicted probability for the token.
      ⭐
      Before version v0.7.0 (see v0.6.4 code), the loss was computed simply as:
      Here, self.loss_fn was a standard cross-entropy loss:
      Average over all non-padding tokens . In practice, the old code masked out prompt tokens by setting their labels to a special ignore index (e.g., -100), so only response tokens contributed.
      Because cross-entropy with an ignore-index behaves exactly like the current version (masked mean of negative log-probabilities), the old implementation is mathematically equivalent to the current SFTLoss formulation.
      Prev
      8.2 Instruction Fine-Tuning Qwen3-0.6B: A Minimal Working Example
      Next
      8.4 OpenRLHF: Reward Model Training and Code Analysis
      Loading...
      Article List
      LLM Learning Roadmap
      ✨ Awesome-Anything
      Lazy loaded image Awesome-AI-Tutorials
      Lazy loaded image Awesome-AI-Researchs
      Lazy loaded image Awesome-GitHub-Repos
      Lazy loaded image Awesome-AI-Examples
      🖼️ Digital Image Processing
      Lazy loaded image DIP I Image Demosaicing and Histogram Manipulation
      Lazy loaded image DIP II Image Denoising
      Lazy loaded image DIP III Edge Detection
      Lazy loaded image DIP IV Digital Half-toning
      Lazy loaded image DIP V Geometric Image Modification
      🍃 LLM Components
      Lazy loaded image 1.1 Tokenization
      Lazy loaded image 1.2 Embedding
      Lazy loaded image 1.3 Attention
      Lazy loaded image 1.4 FFN, Residual Addition, and LN
      Lazy loaded image 1.5 Positional Encoding
      Lazy loaded image 1.6 Architecture & Decoding Policy
      🌱 LLM Pre-training
      Lazy loaded image 2.1 Pre-training Definition
      Lazy loaded image 2.2 Pre-training Data
      Lazy loaded image 2.3 Pre-training Workflow
      Lazy loaded image 2.4 Continued Pre-training (CPT)
      Lazy loaded image 2.5 Pre-training Evaluation
      ☘️ LLM Post-Training
      Lazy loaded image 3.1 Supervised Fine-Tuning (SFT)
      Lazy loaded image 3.2 Reinforcement Learning
      Lazy loaded image 3.3 RLHF: Reinforcement Learning from Human Feedback
      Lazy loaded image 3.4 Optimizing PPO-Based Algorithms
      Lazy loaded image 3.5 Direct Preference Optimization (DPO)
      🍀 LLM Popular Models
      Lazy loaded image 4.1 BERT & Variants
      Lazy loaded image 4.2 PaLM Series
      Lazy loaded image 4.3 GPT Series
      Lazy loaded image 4.4 LLaMA Series
      Lazy loaded image 4.5 GLM Series
      Lazy loaded image 4.6 Qwen Series
      Lazy loaded image 4.7 DeepSeek Series
      Lazy loaded image 4.8 MOE Series
      Lazy loaded image 4.9 Other
      🪴 LLM Applications
      Lazy loaded image 5.1 Prompt Engineering
      Lazy loaded image 5.2 LLM-based Agents
      Lazy loaded image 5.3 Retrieval-Augmented Generation (RAG)
      Lazy loaded image 5.4 Enterprise-level RAG Practice (I)
      Lazy loaded image 5.5 Enterprise-level RAG Practice (II)
      🌿 LLM Optimization
      Lazy loaded image 6.1 GPU Usage Analysis
      Lazy loaded image 6.2 Latency Analysis for Inference
      Lazy loaded image 6.3 FlashAttention
      Lazy loaded image 6.4 PagedAttention
      Lazy loaded image 6.5 Packing for Efficient Training
      Lazy loaded image 6.6 PEFT
      🌾 LLM Compression
      Lazy loaded image 7.1 Quantization
      Lazy loaded image 7.2 Pruning
      Lazy loaded image 7.3 Knowledge Distillation
      🌵 LLM Hands-on Practice
      Lazy loaded image 8.1 nanoGPT: Training a GPT from Scratch
      Lazy loaded image 8.2 Instruction Fine-Tuning Qwen3-0.6B: A Minimal Working Example
      Lazy loaded image 8.3 OpenRLHF: SFT Qwen3-0.6B and Code Analysis
      Lazy loaded image 8.4 OpenRLHF: Reward Model Training and Code Analysis
      Lazy loaded image 8.5 OpenRLHF: PPO Tutorial and Code Analysis
      🌴 LLM Must-read Papers
      Lazy loaded image 9.1 Technical Reports Explained
      Lazy loaded image 9.2 MedLLM
      Lazy loaded image 9.3 o1 Techinical Paper
      Lazy loaded image 9.4 Prompting
      Lazy loaded image 9.5 LLM-as-Judge
      🌳 LLM Q&A
      Lazy loaded image 10.1 Why GRPO training loss starts at 0
      Lazy loaded image 10.2 VLM-R1 Reproduction Results
      Lazy loaded image 10.3 KL(P||Q) or KL(Q||P)?
      Lazy loaded image 10.4 Left Padding or Right Padding?
      Lazy loaded image 10.5 What is QK-Norm?
      Lazy loaded image 10.6 Why LLM training rarely adopts Dropout?
      Lazy loaded image 10.7 Dropout and DropPath
      🐝 VLM Image Encoders
      Lazy loaded image 11.0 Introduction
      Lazy loaded image 11.1 Vision Transformer
      Lazy loaded image 11.2 CLIP
      Lazy loaded image 11.3 BLIP
      Lazy loaded image 11.4 EVA
      Lazy loaded image 11.6 BLIP-2
      Lazy loaded image 11.7 SigLIP
      📝 MISC.
      🤔 Large Language Models Pass the Turing Test.
      Lazy loaded image OpenAI ChatGPT Progression
      ⏳ Claude Sonnet 4 Test with Sandtris
      🔍 What DeepSeek-R1's Surge Teaches Us About SEO
      💰 CivitAI’s Payment Issue
      🎓 arXivSub, an Academia Assistant