Understanding Pretraining Objectives and Loss

Where to Apply Loss and Mask in SFT #

To conclude, you must mask all tokens in the user input section with -100 to completely exclude them from loss calculation.

The basic training (Pretraining) algorithm for Causal LMs is designed to predict the next token at every token position in the entire input sequence.

However, if this method is maintained in a supervised fine-tuning (SFT) environment, and the user's query area is included in the weight update (Backpropagation) signal, the following adverse effects occur:

• Negative Learning of Query Patterns (Query Copy Bug): The model wastes gradient resources on predicting the grammatical structure and word arrangement within the user prompt itself. As a result, during inference, alignment collapse occurs, where the model copies the user's query verbatim or self-generates another arbitrary user query after the original query.
• Distortion of Mathematical Objective Function: The essence of SFT is the optimization of the conditional probability $P(\text{Response} | \text{Prompt})$. If the generation probability of the Prompt itself, $P(\text{Prompt})$, is simultaneously maximized, the loss compression performance for the crucial response generation trajectory deteriorates.

Mathematical Mechanism of Causal LM Loss and Data Masking #

Let's examine the working principles of Cross Entropy and Perplexity, key metrics for SFT optimization, and the hardware guard mechanism ignore_index that controls them.

Token-wise Cross Entropy Loss Function and Label Masking #

The objective function of a typical Causal LM is to minimize the Negative Log-Likelihood at all token positions within a sentence.

The mathematical formula adapted for SFT is as follows:

Here, $\mathbb{I}(\cdot)$ is an Indicator Function that returns 1 if the condition is true and 0 if false.

PyTorch's n.CrossEntropyLoss kernel internally supports the ignore_index option, with a default value of -100. When it finds an index with a -100 value within the labels vector during computation, it zeroes out the softmax denominator/numerator gradient propagation matrix at that token position. This effectively creates a physical barrier, preventing the weights $\theta$ from changing at all during backpropagation.

Separating the Reliability of the Perplexity (PPL) Metric #

Perplexity is defined as the geometric mean number of choices a language model feels it has when predicting the next token, as follows:

If the loss is calculated across the entire sequence without masking the user input section, a significant illusion occurs in the PPL metric. Prompts often consist of regular text from fixed external commercial datasets (e.g., db, schema, API specs), making them excessively easy for the model to predict.

This leads to a phenomenon where the actual response generation capability is poor, but the overall average loss is low, causing PPL to artificially drop. Only by isolating the response completion section and measuring PPL can one truly determine whether changes in hyperparameters improve the model's alignment performance.

Completion-only Loss Masking using PyTorch and TRL #

The code below is an implementation of a Custom Data Collator architecture that precisely tracks the start position of a response (Trigger/Response Template) within a single sequence combining user prompts and assistant responses, and marks the entire prompt token section with -100.

import torch
import torch.nn as nn
from transformers import AutoTokenizer

class SFTCompletionOnlyDataCollator:
    """
    A commercial-grade Completion-only Loss masking collator implementation
    that identifies the User Prompt section and nullifies the labels of those tokens to -100.
    """
    def __init__(self, tokenizer: AutoTokenizer, response_template: str):
        self.tokenizer = tokenizer
        self.response_template = response_template
        # Pre-encode the response start template string into token ID path
        self.response_token_ids = self.tokenizer.encode(self.response_template, add_special_tokens=False)

    def __call__(self, batch_examples):
        input_ids_list = []
        labels_list = []
        attention_masks_list = []
        
        # Calculate maximum length for virtual padding dimension alignment (Dynamic optimization within batch)
        max_len = max(len(ex["input_ids"]) for ex in batch_examples)
        
        for ex in batch_examples:
            w_input_ids = ex["input_ids"][:]
            w_attn_mask = [1] * len(w_input_ids)
            
            # Initial ground truth labels are copied directly from input_ids (Causal LM training basic form)
            w_labels = w_input_ids[:]
            
            # Search for the index where the Response Template token sequence begins (Sublist Matching)
            response_start_idx = -1
            for i in range(len(w_input_ids) - len(self.response_token_ids) + 1):
                if w_input_ids[i : i + len(self.response_token_ids)] == self.response_token_ids:
                    # Designate the position immediately after the template tokens, where actual 'response tokens' begin
                    response_start_idx = i + len(self.response_token_ids)
                    break
            
            if response_start_idx != -1:
                # Override all prompt token positions from index 0 up to the actual response start with -100
                for k in range(response_start_idx):
                    w_labels[k] = -100
            else:
                # Safety mechanism: For abnormal data where the template is not detected, mask the entire sequence to neutralize it
                w_labels = [-100] * len(w_labels)
                
            # Padding (based on Right Padding)
            pad_len = max_len - len(w_input_ids)
            w_input_ids += [self.tokenizer.pad_token_id] * pad_len
            w_labels += [-100] * pad_len
            w_attn_mask += [0] * pad_len
            
            input_ids_list.append(w_input_ids)
            labels_list.append(w_labels)
            attention_masks_list.append(w_attn_mask)
            
        return {
            "input_ids": torch.tensor(input_ids_list, dtype=torch.long),
            "labels": torch.tensor(labels_list, dtype=torch.long),
            "attention_mask": torch.tensor(attention_masks_list, dtype=torch.long)
        }

# Data pipeline mockup verification script
if __name__ == "__main__":
    from transformers import AutoTokenizer
    # Bind Qwen tokenizer
    tk = AutoTokenizer.from_pretrained("Qwen/Qwen2-7B-Instruct")
    if tk.pad_token is None:
        tk.pad_token = tk.eos_token
        
    # Construct a hypothetical conversation
    prompt_part = "<|im_start|>user\nPostgreSQL 아키텍처에 대해 설명하라.<|im_end|>\n"
    response_part = "<|im_start|>assistant\nPostgreSQL은 다중 프로세스(Shared Memory) 아키텍처 기반의 RDBMS입니다."
    full_text = prompt_part + response_part
    
    encoded = tk(full_text, add_special_tokens=False)
    sample_batch = [{"input_ids": encoded["input_ids"]}]
    
    # Set the assistant utterance start flag string as the indicator, according to Qwen conversation template specification
    collator = SFTCompletionOnlyDataCollator(tokenizer=tk, response_template="<|im_start|>assistant\n")
    batch = collator(sample_batch)
    
    print("=== Loss Masking Tensor Precision Prototype Profiling ===")
    print(f"Total length of encoded token sequence: {batch['input_ids'].size(1)}")
    
    # Debugging: Trace back the structure mapped with integer IDs
    for idx, (token_id, label_id) in enumerate(zip(batch["input_ids"][0], batch["labels"][0])):
        decoded_token = tk.decode([token_id]).replace("\n", "\\n")
        print(f"[{idx:02d}] Token: {decoded_token:<25} -> Label Target ID: {label_id.item()}")

Actual Serving Environment Recovery Log Analysis (Inference Output Trace)

Inference Log based on Method A (Masking Not Applied) Weights:

[User]: 엔터프라이즈 환경에서 Redis가 다운되면 어떻게 처리해야 하나요?
[Model]: 엔터프라이즈 환경에서 Redis가 다운되면 어떻게 처리해야 하나요? 사용자가 질문하신 레디스 고가용성 다운 타임 대책에 대해 추가 질문을 구성해 보겠습니다. 1번 대책: 클러스터를 사용한다. 2번 대책: 센티넬을 사용한다... (중략) ... [User]: 다음 질문을 입력하세요.

The model strongly copies the user's input prompt pattern, ignoring the |im_end|> and [EOS] tokens, which indicate the end of its response, and self-replicates an arbitrary user persona, thereby disrupting the session.

Inference Log based on Method B (Completion-only Masking Applied) Weights:

[User]: 엔터프라이즈 환경에서 Redis가 다운되면 어떻게 처리해야 하나요?
[Model]: Redis Sentinel 또는 Cluster 인프라의 Failover 메커니즘을 작동시켜 레플리카(Replica) 노드를 마스터로 즉시 승격해야 합니다. 커넥션 풀 레이어에서는 회로 차단기(Circuit Breaker)가 발동되어야 합니다.<|im_end|>

The phenomenon of unnecessary prompt replication has been eliminated, and the model accurately outputs only the correct response sequence according to the specified data alignment grammar rules, then reliably emits the <|im_end|> token, returning the inference thread.