Lazy loaded image6.5 Packing for Efficient Training

6.5 Packing for Efficient Training

💡
Why
In traditional model training, a fixed-length input sequence is typically used. However, the lengths of text segments in real-world data vary greatly. If short texts are directly padded to a fixed length using [PAD] tokens, it results in a large amount of wasted computation.
How
Packing refers to concatenating multiple short sequences into one long sequence, allowing each batch of tokens to be closer to the maximum length limit and maximizing token utilization. In practice, packing can be done by filling up to the longest sentence in a batch or to the model's maximum input length.

6.5.1 Pre-training Phase

📌
Most packing during pretraining is simply concatenation, where texts from different sources are separated by special tokens, such as the special token [SEP].
Of course, if a long text is accidentally split mid-way, it could cause issues during the pre-training stage. In fact, we can tolerate some longer texts being split during initial pretraining (not during continued pretraining on long texts), but if this splitting is too aggressive, how can we solve it?
A good approach is to pack short and long texts separately. Short and long texts can each have their own packing processes, with pack sizes matching their original max length (e.g., 4k for short texts, 128k for long ones). Then, repack the packed short texts again to longer ones, i.e., 128k.

6.5.2 Supervised Fine-tuning Phase

🎄
An Efficient Method Without Packing
The efficiency here is mainly applied to multi-turn dialogue scenarios, which usually do not need packing because we naturally put all-turn dialogue together to prevent computation waste. But we need to make sure the loss is calculated correctly.
Since causal language modeling (CLM) uses a triangle attention mask, the structure ensures that each token can only attend to itself and the tokens before it. Therefore, when predicting a response token, the model can access all prior information up to that token, but anything after it remains inaccessible.
Based on this mechanism, we can precompute logits for all responses in a dialogue at once, and when calculating loss, retain only the response portion of the logits—ignoring the user prompt portion—to compute loss correctly and more efficiently.
🛠
Let’s look at how Llama Factory improves efficiency
Thanks to current transformers supporting the 4D mask mechanism, the latest implementations now differentiate attention masks. This allows marking of different document types or sequences and proper handling of padding regions (usually 0) with their own identifiers.
📌
What is a 4D mask?
Standard attention masks are given in 2D tensors, shaped [batch_size, sequence_length]. Inside the LLM model, 2D tensors are expanded into 4D tensors, shaped as [batch_size, heads, sequence_length, sequence_length].
If we can pass our own mask, this format would allow for more fine-grained attention strategies. For example, in causal decoding, it uses a lower triangular matrix with ones like the left one of below.
notion image
To get a separate mask for two different sequences, we need to custom our own mask. The so-called 4D Mask supports custom attention mask inputs.
For example, let's say we have a prefix with The Last Of, and by applying beam search, we have 3 beams with token like You, Me, and Us. The standard approach would make them in a batch of 3 samples, with the shape of (3, 4):
With 4D Mask, we would have just one batch shaped (1, 6):
Then the custom 4D mask would have a shape of (1, 1, 6, 6):
And the positions_ids would be
In this way, you can observe that we can generate whatever 4D masks we want, like non-sequential tokens (like the example above) or multiple sequences in a single input (like packing multiple short texts).
For more details, refer to the original implementation about 4D Mask on GitHub or the author’s post.
<ins/>
  • Example code (Running on Colab recommended)
      • you may see results like
Prev
6.4 PagedAttention
Next
6.6 PEFT
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