Sequence Packing and Loss Masking in Production SFT: Architecture, Trade-Offs, and Framework Implementations

Supervised fine-tuning pipelines routinely waste 40–70% of FLOPs on padding tokens. Sequence packing eliminates that waste by concatenating multiple training examples into a single long tensor, but it introduces three hard constraints: attention must not leak across example boundaries, loss must be computed only on the intended target spans (typically assistant responses), and the data collation logic must stay compatible with Flash Attention 2 kernels. This post surveys how the major open-source frameworks — Hugging Face TRL, Axolotl, LLaMA-Factory, Unsloth, and NVIDIA NeMo — solve these constraints, where their designs converge, and where they still diverge.

Why Packing Changes the Loss Masking Problem

In a padded batch, every example has the same length. A binary attention_mask (1 for real tokens, 0 for padding) is sufficient: the loss function simply ignores positions where attention_mask == 0 via CrossEntropyLoss(ignore_index=-100). When examples are packed, a single tensor contains several independent sequences separated only by an EOS token (or nothing at all). The same ignore_index = -100 mechanism still works for padding, but now two distinct masking jobs coexist:

• Padding mask — Exclude pad tokens introduced by the collator. Label value -100.
• Prompt/span mask — Exclude system, user, and tool-result spans so loss only covers assistant responses. Also labeled -100 (same index, different semantic meaning).

Because both jobs reuse ignore_index, a single label tensor can express both — if the preprocessing pipeline produces correct span boundaries for every packed example. The difficulty is that chat templates turn a multi-turn conversation into a flat token stream; the span boundaries (where user ends and assistant begins) must be computed before packing, then preserved through concatenation.

Attention Isolation: Block-Diagonal vs. Variable-Length Kernels

Two architectural families exist for preventing cross-example attention in a packed sequence.

1. Block-Diagonal Attention Masks (4D masks)

The conventional approach builds a (batch, 1, seq_len, seq_len) boolean mask where each example's block is causal and off-block entries are False. Complexity grows as (sum s_i)^2 instead of sum s_i^2, so practical packing length is limited. Axolotl's non-Flash path and early TRL packing used this method.

2. Variable-Length Flash Attention (flash_attn_varlen_func)

Flash Attention 2 accepts a 1D cu_seqlens tensor (cumulative sequence lengths) instead of a square mask. Attention is computed only within each [cu_seqlens[i], cu_seqlens[i+1]) window. Complexity stays at sum s_i^2, enabling arbitrary packing lengths. The collator must emit:

• Flattened input_ids, labels, position_ids (reset to 0 at each sequence start)
• cu_seqlens for the varlen kernel
• max_seqlen for kernel launch bounds

Hugging Face's DataCollatorWithFlattening, Axolotl's Multipack, Unsloth's padding-free batching, and NeMo's sequence packing all adopt this path.

Framework-by-Framework Comparison

Hugging Face TRL — Packing strategy: DataCollatorForCompletionOnlyLM(padding_free=True) + SFTTrainer. Loss masking: response_template string match; train_on_inputs=False masks everything before the template. Flash Attention integration: DataCollatorWithFlattening emits cu_seqlens + position_ids. Config surface: padding_free=True flag on collator.

Axolotl — Packing strategy: multipack: true (Flash) or 4D mask (non-Flash). Loss masking: roles_to_train: ["assistant"], train_on_eos: "turn", chat_template spans. Flash Attention integration: native cu_seqlens emission in collator. Config surface: multipack: true, flash_attention: true.

LLaMA-Factory — Packing strategy: packing: true (wrapped) + neat_packing: true (FFD, SFT only). Loss masking: train_on_prompt: false, mask_history: true (last turn only). Flash Attention integration: HF DataCollatorForSeq2Seq with position_ids; FFD uses TRL 0.19+ strategy. Config surface: packing, neat_packing, train_on_prompt, mask_history.

Unsloth — Packing strategy: packing=True in SFTConfig + padding-free batching. Loss masking: inherits TRL collator behavior; train_on_inputs via DataCollatorForCompletionOnlyLM. Flash Attention integration: custom Triton kernels (Fused QK RoPE + packing) reset RoPE position IDs per sequence. Config surface: packing=True, max_length.

NVIDIA NeMo — Packing strategy: offline packing script (prepare_packed_ft_dataset.py) produces .npy bins. Loss masking: GPTSFTDataset builds loss masks during tokenization; spans from chat template. Flash Attention integration: TransformerEngine / FlashAttention varlen kernels; packed_sequence=True. Config surface: packed_sequence, pack_sizes, packing_algorithm.

Key Convergences

• All production paths now require Flash Attention 2 (or TransformerEngine equivalent) for efficient packing.
• All expose position IDs reset per sequence as the mechanism that lets varlen kernels work.
• All reuse ignore_index=-100 for both padding and prompt masking — the distinction is purely in how the label tensor is constructed upstream.

Key Divergences

• When masking happens — Axolotl and NeMo do it during dataset preprocessing (offline); TRL and Unsloth do it in the collator (online); LLaMA-Factory splits the difference.
• Multi-turn granularity — Axolotl's train_on_eos: "turn" masks per assistant turn; LLaMA-Factory's mask_history=true keeps only the last turn; TRL's response_template matches once per example (single-turn assumption unless you pre-split).
• Packing algorithm — NeMo uses offline first_fit_shuffle / first_fit_decreasing; Axolotl uses an online multi-pack sampler; TRL 0.19+ uses First-Fit Decreasing (FFD); Unsloth uses custom bin-packing with length metadata.
• MoE / expert parallelism — NeMo and Axolotl document interaction with expert parallelism; TRL and Unsloth are less explicit.

Multi-Turn Loss Masking: Three Common Policies

Given a conversation [sys, u1, a1, u2, a2, u3, a3], frameworks let you choose:

1. Train on all assistant turns — default in Axolotl (roles_to_train: ["assistant"]), TRL (if you split into per-turn examples), NeMo.
2. Train only on the final turn — LLaMA-Factory mask_history=true; Axolotl train_on_eos: "last".
3. Weighted mixing — assign lower weight to earlier turns. Sebastian Raschka notes this requires unreduced cross-entropy plus custom reduction; no framework exposes it natively yet.

Pitfall: Role markers like <|im_start|>assistant often tokenize to multiple IDs. String-search masking ("assistant" in decoded text) fails. All frameworks now derive spans from the chat template's token offsets — Axolotl via tokenizer.apply_chat_template(..., return_dict=True), TRL via DataCollatorForCompletionOnlyLM's template-aware tokenizer pass, NeMo in GPTSFTDataset.

Practical Checklist for a Packed SFT Run

1. Verify model supports position_ids — Llama 2/3, Mistral, Mixtral, Granite, DBRX, Falcon, Gemma, OLMo, Phi 1/2/3, Qwen 2/2-MoE, StableLM, StarCoder 2 (per the HF blog).
2. Enable Flash Attention 2 — attn_implementation="flash_attention_2" in model load.
3. Pick one packing path — do not mix dynamic batching and packing; NeMo-RL warns they are exclusive.
4. Inspect one batch before training — decode input_ids, labels, position_ids, cu_seqlens; confirm position_ids resets to 0 at each sequence start, labels == -100 exactly on prompt/pad spans, and cu_seqlens length equals num_sequences_in_pack + 1.
5. Adjust global batch size — each pack contains approximately pack_size / avg_seq_len examples; divide the global batch size by that factor to keep tokens-per-step constant.
6. Monitor throughput (tokens/s) and validation loss — packing should not degrade convergence. The HF blog shows identical validation loss; NeMo claims no impact.

Open Gaps

• Weighted multi-turn loss (intermediate between "all turns" and "last turn only") requires custom loss reduction; no YAML flag exists in any framework.
• Tool-call masking — environment-provided tool results should be masked like user turns, but only Axolotl's content-parts format (reasoning_content field) handles this explicitly.
• Cross-framework packing format portability — NeMo's .npy bins, Axolotl's online sampler, and TRL's online collator produce incompatible artifacts. No standard interchange format exists.
• Evaluation during packed training — validation sets are usually not packed, so throughput numbers are not directly comparable.

Sources

• Hugging Face Blog: "Improving Hugging Face Training Efficiency Through Packing with Flash Attention 2" — https://huggingface.co/blog/packing-with-FA2
• Sebastian Raschka: "When to mask prompt tokens during SFT" — https://sebastianraschka.com/faq/docs/when-mask-prompt-tokens.html
• NVIDIA NeMo Framework User Guide: "Sequence Packing" — https://docs.nvidia.com/nemo-framework/user-guide/24.12/nemotoolkit/features/optimizations/sequence_packing.html
• Axolotl Documentation: "Multipack (Sample Packing)" — https://docs.axolotl.ai/docs/multipack.html
• Axolotl Documentation: "Conversation Dataset Format" — https://docs.axolotl.ai/docs/dataset-formats/conversation.html
• LLaMA-Factory Arguments Reference — https://llamafactory.readthedocs.io/en/latest/advanced/arguments.html
• Unsloth Blog: "3x Faster LLM Training with Unsloth Kernels + Packing" — https://unsloth.ai/docs/blog/3x-faster-training-packing
• ArXiv 2407.09105: "Enhancing Training Efficiency Using Packing with Flash Attention" — https://arxiv.org/html/2407.09105v1
• Kaitchup Substack: "Padding-Free vs. Packing: Fast and Efficient Fine-Tuning for LLMs Explained" — https://kaitchup.substack.com/p/padding-free-vs-packing-fast-and
• Together AI Blog: "Fine-Tuning LLMs for Multi-Turn Conversations: A Technical Deep Dive" — https://www.together.ai/blog/fine-tuning-llms-for-multi-turn-conversations-a-technical-deep-dive
• Hugging Face TRL Issue #805: "Packing in SFT" — https://github.com/huggingface/trl/issues/805