SPDA, FA2 vs. Eager Attention Implementation leading to different losses

[jiosephlee]

System Info

Hi this is using TRL but it seems like a lower level issue.

I'm training a variant of Qwen3 (Intern-S1-mini) but I'm not using the vision tower so it's effectively Qwen3-8B. I've been doing finetuning and checking different attention implementations i.e. SPDA vs. Flash Attention 2. However, I've been getting strange results where the downstream test accuracy is different (FA2 is worse). Furthermore, it seems like this issue is accentuated with grad accumulation. I'm not sure what's the best way to share this as my current code abstracts upon HF Trainer for my personal convenience.

Who can help?

No response

Information

• The official example scripts
• My own modified scripts

Tasks

• An officially supported task in the examples folder (such as GLUE/SQuAD, ...)
• My own task or dataset (give details below)

Reproduction

Here are the current values of my config

"per_device_train_batch_size": 16,
"gradient_accumulation_steps": 2,
"optim": "paged_adamw_8bit",
"evaluation_strategy": "epoch",
"weight_decay": 0.1,
"gradient_checkpointing": true,
"use_liger_kernel": true,
"num_train_epochs": 1,
"learning_rate": 8e-05,
"lr_scheduler_type": "cosine",
"warmup_steps": 0,
"warmup_ratio": 0.1,
"report_to": "wandb",
"run_name": "finetune_Tox_internlm_Intern-S1-mini",
"logging_steps": 1,
"logging_strategy": "steps",
"save_strategy": "no",
"remove_unused_columns": false,
"seed": 42,
"completion_only_loss": false,
"dataset_text_field": "text",
"packing": false,
"padding_free": false,
"loss_type": "nll"```

### Expected behavior

They should have equal test accuracy.

[Rocketknight1]

Hmm, this is quite hard for us to debug. You might need to explore a bit more - if you find that one of the attention backends is giving incorrect outputs, that's definitely something we should investigate, but we can't really do much without some clear code that indicates the exact issue. The more specific you can be, the more likely it is that we'll be able to fix it!

[jiosephlee]

@Rocketknight1 Here's a script that is self-contained as well as the loss charts.

import json
import torch
from transformers import AutoModelForCausalLM, AutoTokenizer
from datasets import load_dataset
from trl import SFTTrainer, SFTConfig
import logging
from typing import List
import wandb

# Set up basic logging
logging.basicConfig(
    level=logging.INFO,
    format="%(asctime)s - %(levelname)s - %(message)s",
    datefmt="%Y-%m-%d %H:%M:%S",
)
log = logging.getLogger(__name__)

def apply_prompt_template(examples):
    """Create a single 'text' field with prompt+completion for SFTTrainer."""
    prompts = [f"Review: {t}\nSentiment:" for t in examples["text"]]
    completions = [" Positive" if lbl == 1 else " Negative" for lbl in examples["label"]]
    return {"text": [p + c for p, c in zip(prompts, completions)]}

def run_benchmark(config):
    """
    Runs a single benchmark configuration.
    """
    log.info(f"--- Starting Benchmark Run: {config['run_name']} ---")

    # --- 1. Load Model and Tokenizer ---
    log.info(f"Loading model: {config['model_id']} with attn_implementation: {config['attn_implementation']}")
    model = AutoModelForCausalLM.from_pretrained(
        config['model_id'],
        attn_implementation=config['attn_implementation'],
        dtype=torch.bfloat16,
        trust_remote_code=True,
        device_map="cuda"
    )
    tokenizer = AutoTokenizer.from_pretrained(config['model_id'], trust_remote_code=True)
    if tokenizer.pad_token is None:
        tokenizer.pad_token = tokenizer.eos_token
        model.config.pad_token_id = model.config.eos_token_id

    # --- 2. Load and Prepare Dataset (single 'text' field) ---
    log.info(f"Loading and preparing dataset: {config['dataset_name']}")
    dataset = load_dataset(config['dataset_name'])
    
    train_dataset = dataset['train'].shuffle(seed=42).select(range(100))
    test_dataset = dataset['test'].shuffle(seed=42).select(range(500))
    
    # Map to a single 'text' column; SFTTrainer will tokenize
    train_text_dataset = train_dataset.map(apply_prompt_template, batched=True)
    test_text_dataset = test_dataset.map(apply_prompt_template, batched=True)
    train_text_dataset = train_text_dataset.remove_columns([c for c in train_text_dataset.column_names if c != "text"])
    test_text_dataset = test_text_dataset.remove_columns([c for c in test_text_dataset.column_names if c != "text"])

    # --- 3. Configure Training ---
    training_output_dir = os.path.join(config['output_dir'], config['run_name'])

    training_args = SFTConfig(
        run_name=config['run_name'],
        dataset_text_field = "text",
        num_train_epochs=config['epochs'],
        per_device_train_batch_size=config['per_device_batch_size'],
        gradient_accumulation_steps=config['gradient_accumulation_steps'],
        learning_rate=config['learning_rate'],
        logging_dir=f"{training_output_dir}/logs",
        gradient_checkpointing=False,
        logging_steps=1,
        bf16=True,
        report_to="wandb",
        remove_unused_columns=False,
    )

    trainer = SFTTrainer(
        model=model,
        args=training_args,
        train_dataset=train_text_dataset,
        processing_class=tokenizer,
    )

    # --- 4. Train ---
    log.info("Starting training...")
    trainer.train()

    # --- 5. Custom Evaluation: first-token logits for " Positive" vs " Negative" ---
    log.info("Starting custom evaluation (first-token logits)...")
    device = model.get_input_embeddings().weight.device

    # Candidate first-token ids (with leading space)
    pos_tokens: List[int] = tokenizer.encode(" Positive", add_special_tokens=False)
    neg_tokens: List[int] = tokenizer.encode(" Negative", add_special_tokens=False)
    if len(pos_tokens) == 0 or len(neg_tokens) == 0:
        raise ValueError("Tokenization of candidates returned empty sequences.")
    pos_first_id, neg_first_id = pos_tokens[0], neg_tokens[0]

    # Build prompts (without answers)
    prompts = [f"Review: {t}\nSentiment:" for t in test_dataset["text"]]

    batch_size = config.get("per_device_batch_size")
    num_batches = (len(prompts) + batch_size - 1) // batch_size

    correct = 0
    total = 0
    margins: List[float] = []

    model.eval()
    with torch.no_grad():
        for i in range(num_batches):
            batch_prompts = prompts[i * batch_size:(i + 1) * batch_size]
            batch_labels = test_dataset["label"][i * batch_size:(i + 1) * batch_size]

            tokenized = tokenizer(
                batch_prompts,
                return_tensors="pt",
                padding=True,
                truncation=True,
                
                max_length=512,
            )
            input_ids = tokenized["input_ids"].to(device)
            attention_mask = tokenized["attention_mask"].to(device)

            outputs = model(input_ids=input_ids, attention_mask=attention_mask, use_cache=False)
            logits = outputs.logits  # [B, T, V]

            lengths = attention_mask.sum(dim=1) - 1  # index of last prompt token
            batch_indices = torch.arange(logits.size(0), device=device)
            next_token_logits = logits[batch_indices, lengths, :]  # [B, V]

            pos_logits = next_token_logits[:, pos_first_id]
            neg_logits = next_token_logits[:, neg_first_id]
            preds = (pos_logits > neg_logits).long().tolist()  # 1 for Positive, 0 for Negative

            # Update metrics
            for pred, label, pos_l, neg_l in zip(preds, batch_labels, pos_logits.tolist(), neg_logits.tolist()):
                correct += int(pred == int(label))
                total += 1
                margins.append(float(pos_l - neg_l))

    accuracy = correct / max(total, 1)
    avg_margin = float(sum(margins) / max(len(margins), 1))
    log.info(f"Custom eval accuracy (first-token): {accuracy:.4f}, avg logit margin (pos-neg): {avg_margin:.4f}")

    # --- 6. Save Results ---
    results = {
        "config": config,
        "custom_eval": {
            "accuracy_first_token": accuracy,
            "avg_logit_margin_pos_minus_neg": avg_margin,
        },
    }
    
    results_path = os.path.join(config['output_dir'], f"{config['run_name']}_results.json")
    os.makedirs(config['output_dir'], exist_ok=True)
    with open(results_path, 'w') as f:
        json.dump(results, f, indent=4)
        
    log.info(f"Results for {config['run_name']} saved to {results_path}")

def main():
    # --- Define Benchmark Configurations ---
    base_config = {
        "model_id": "Qwen/Qwen2.5-1.5B",
        "dataset_name": "imdb",
        "epochs": 1,
        "per_device_batch_size": 2,
        "learning_rate": 5e-5,
        "output_dir": "./results_benchmark"
    }

    configurations = [
        {**base_config, "attn_implementation": "eager", "gradient_accumulation_steps": 4, "run_name": "eager_with_grad_accum"},
        {**base_config, "attn_implementation": "eager", "gradient_accumulation_steps": 1, "run_name": "eager_no_grad_accum"},
        {**base_config, "attn_implementation": "flash_attention_2", "gradient_accumulation_steps": 4, "run_name": "flash_attention_2_with_grad_accum"},
        {**base_config, "attn_implementation": "flash_attention_2", "gradient_accumulation_steps": 1, "run_name": "flash_attention_2_no_grad_accum"},
    ]

    for config in configurations:
        run_benchmark(config)
        wandb.finish()

if __name__ == "__main__":
    main()

The issue is between eager attention and any other form of attention (e.g. SPDA, FlashAttention2). While the variance in loss for this script is minimal, I'm getting differences in test accuracy as large as 5% for larger and more complex tasks. The issue also seems worse with grad accumulation, but I'm less sure of this. In terms of why this happens, it could be a precision issue https://arxiv.org/abs/2405.02803

[Rocketknight1]

Hmmn, this difference seems within the limits of what we'd expect just from small numerical changes! Different attention implementations, even when mathematically equivalent, can give very slightly different outputs because different orders of floating-point operations can yield slightly different numerical errors.

[jiosephlee]

Ah, I see. It also seems like I just got a few bad datasets where eager attention did better. Upon further exploration, the variance observed in downstream accuracy for larger tasks doesn't seem actually biased towards any attention implementation. Thanks!