[SFT] High vRAM consumption during eval loop

[Khreas]

Reproduction

Unexpected behavior

When training a model on large sequences (>=20k tokens) with PEFT LoRA + SFTTrainer + liger-kernel, the vRAM usage spikes during the evaluation loop, consuming way more vRAM than during the training.

The size of this vRAM spike seem to scale with the length of the input sequence: for cases with max_length=40000, we end up with spikes of ~50GB vRAM, far exceeding the amount used during the training.

Here's a MLFlow GPU vRAM extract showcasing this on an A100 for this 40k token scenario with Qwen3-0.6B:

And same goes for Qwen3-4B, 40k token:

Minimal reproduction script

Below is the default SFT example from the documentation, slightly altered to artificially create long input sequences (>=20k tokens) in both the training and evaluation dataset splits.

By running watch -n 1 nvidia-smi while the training is running, you can see that the vRAM usage is way higher during the evaluation phase than during the training. If your GPU has enough vRAM, you can increase the max_length parameter and this will become even more visible. _For some reason, I can't get trackio to properly report vRAM usage, hence the use of nvidia-smi.

You can launch the script with the following command:

python sft_example.py \
--model_name_or_path Qwen/Qwen3-0.6B \
--dataset_name trl-lib/Capybara \
--learning_rate 2.0e-4  \
--max-steps 10 \
--per_device_train_batch_size 1 \
--per_device_eval_batch_size 1 \
--eval_accumulation_steps 1 \
--gradient_accumulation_steps 1 \
--gradient_checkpointing \
--eos_token '<|im_end|>' \
--eval_strategy steps \
--eval_steps 10 \
--use_peft \
--lora_r 8 \
--lora_alpha 16 \
--use_liger \
--max_length 10000

# Copyright 2020-2025 The HuggingFace Team. All rights reserved.
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
#     http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.

# /// script
# dependencies = [
#     "trl",
#     "peft",
#     "trackio",
#     "kernels"
# ]
# ///

import argparse
import os

from accelerate import logging
from datasets import load_dataset
from transformers import AutoConfig, AutoModelForCausalLM
from transformers.models.auto.modeling_auto import (
    MODEL_FOR_IMAGE_TEXT_TO_TEXT_MAPPING_NAMES,
)
from trl import (
    DatasetMixtureConfig,
    ModelConfig,
    ScriptArguments,
    SFTConfig,
    SFTTrainer,
    TrlParser,
    get_dataset,
    get_kbit_device_map,
    get_peft_config,
    get_quantization_config,
)

logger = logging.get_logger(__name__)

# Enable logging in a Hugging Face Space
os.environ.setdefault("TRACKIO_SPACE_ID", "trl-trackio")


def main(script_args, training_args, model_args, dataset_args):
    ################
    # Model init kwargs
    ################
    model_kwargs = dict(
        revision=model_args.model_revision,
        trust_remote_code=model_args.trust_remote_code,
        attn_implementation=model_args.attn_implementation,
        dtype=model_args.dtype,
    )
    quantization_config = get_quantization_config(model_args)
    if quantization_config is not None:
        # Passing None would not be treated the same as omitting the argument, so we include it only when valid.
        model_kwargs["device_map"] = get_kbit_device_map()
        model_kwargs["quantization_config"] = quantization_config

    # Create model
    config = AutoConfig.from_pretrained(model_args.model_name_or_path)
    valid_image_text_architectures = MODEL_FOR_IMAGE_TEXT_TO_TEXT_MAPPING_NAMES.values()

    if config.architectures and any(
        arch in valid_image_text_architectures for arch in config.architectures
    ):
        from transformers import AutoModelForImageTextToText

        model = AutoModelForImageTextToText.from_pretrained(
            model_args.model_name_or_path, **model_kwargs
        )
    else:
        model = AutoModelForCausalLM.from_pretrained(
            model_args.model_name_or_path, **model_kwargs
        )

    # Load the dataset
    if dataset_args.datasets and script_args.dataset_name:
        logger.warning(
            "Both `datasets` and `dataset_name` are provided. The `datasets` argument will be used to load the "
            "dataset and `dataset_name` will be ignored."
        )
        dataset = get_dataset(dataset_args)
    elif dataset_args.datasets and not script_args.dataset_name:
        dataset = get_dataset(dataset_args)
    elif not dataset_args.datasets and script_args.dataset_name:
        dataset = load_dataset(
            script_args.dataset_name,
            name=script_args.dataset_config,
            streaming=script_args.dataset_streaming,
        )
    else:
        raise ValueError("Either `datasets` or `dataset_name` must be provided.")

    # Emulating long system prompt to analyze vRAM usage
    for split in dataset.keys():
        dataset[split] = (
            dataset[split]
            .map(
                lambda x: {
                    **x,
                    "messages": [
                        {
                            "role": "system",
                            "content": "This string contains 10 tokens exactly. "
                            * 2000,
                        }
                    ]
                    + x["messages"],
                }
            )
            .select(range(10))
        )

    print("Formatted dataset with large system prompt.")

    # Initialize the SFT trainer
    trainer = SFTTrainer(
        model=model,
        args=training_args,
        train_dataset=dataset[script_args.dataset_train_split],
        eval_dataset=dataset[script_args.dataset_test_split]
        if training_args.eval_strategy != "no"
        else None,
        peft_config=get_peft_config(model_args),
    )

    # Train the model
    trainer.train()

    # Log training complete
    trainer.accelerator.print("✅ Training completed.")

    # Save and push to Hub
    trainer.save_model(training_args.output_dir)
    trainer.accelerator.print(f"💾 Model saved to {training_args.output_dir}.")

    if training_args.push_to_hub:
        trainer.push_to_hub(dataset_name=script_args.dataset_name)
        trainer.accelerator.print(
            f"🤗 Model pushed to the Hub in https://huggingface.co/{trainer.hub_model_id}."
        )


def make_parser(subparsers: argparse._SubParsersAction | None = None):
    dataclass_types = (ScriptArguments, SFTConfig, ModelConfig, DatasetMixtureConfig)
    if subparsers is not None:
        parser = subparsers.add_parser(
            "sft", help="Run the SFT training script", dataclass_types=dataclass_types
        )
    else:
        parser = TrlParser(dataclass_types)
    return parser


if __name__ == "__main__":
    print(os.getenv("VIRTUAL_ENV"))
    # from liger_kernel.transformers import AutoLigerKernelForCausalLM  # noqa: F401

    parser = make_parser()
    # When using the trl cli, this script may be run with additional arguments, corresponding accelerate arguments.
    # To ensure that their parsing does not interfere with the script arguments, parse the arguments with
    # `return_remaining_strings=True`, then ignore the remaining strings.
    script_args, training_args, model_args, dataset_args, _ = (
        parser.parse_args_and_config(return_remaining_strings=True)
    )
    main(script_args, training_args, model_args, dataset_args)

System Info

Note: this was also tested on A100 instances.

• Platform: Linux-6.6.87.2-microsoft-standard-WSL2-x86_64-with-glibc2.35
• Python version: 3.12.10
• TRL version: 0.24.0
• PyTorch version: 2.6.0
• accelerator(s): NVIDIA NVIDIA GeForce RTX 4090
• Transformers version: 4.57.3
• Accelerate version: 1.12.0
• Accelerate config: not found
• Datasets version: 4.3.0
• HF Hub version: 0.36.0
• bitsandbytes version: 0.48.2
• DeepSpeed version: not installed
• Liger-Kernel version: 0.6.4
• LLM-Blender version: not installed
• OpenAI version: not installed
• PEFT version: 0.18.0
• vLLM version: not installed

Checklist

• I have checked that my issue isn't already filed (see open issues)
• I have included my system information
• Any code provided is minimal, complete, and reproducible (more on MREs)
• Any code provided is properly formatted in code blocks, (no screenshot, more on code blocks)
• Any traceback provided is complete

[Khreas]

By digging further into the issue, it seems that this vRAM spike comes from the accumulation of the logits during the loss computation. Based on the discussion in this issue on the Liger-kernel repository, the base implementation of the cross-entropy loss creates the following vRAM footprint:

$3 * batch * seq * vocab * fp32_{dtype}$

Which for Qwen3 amounts to:

• 10k tokens

  • $3 * 1 * 10000 * 151936 * 4 = 18GB$

• 40k tokens

  • $3 * 1 * 40000 * 151936 * 4 = 73GB$

This is a bit larger than the vRAM spikes I observed, but it still remains in the same ballpark (I observed $15GB$ for 10k and $50GB$ for 40k).

By looking into the liger-kernel implementation, it seems that the loss used can be either the LigerForCausalLMLoss (which doesn't accumulate logits) and the default self.loss_function defined in the base model used (which accumulate logits):

	# Fused / Chunked loss
    if skip_logits:
        result = LigerForCausalLMLoss(
            hidden_states=kept_hidden_states,
            lm_head_weight=self.lm_head.weight,
            labels=labels,
            shift_labels=shift_labels,
            hidden_size=self.config.hidden_size,
            **kwargs,
        )
        loss, _, token_accuracy = unpack_cross_entropy_result(result)

	
    # Default loss from HF
    else:
        logits = self.lm_head(kept_hidden_states)
        if labels is not None or shift_labels is not None:
            loss = self.loss_function(
                logits=logits,
                labels=labels,
                shift_labels=shift_labels,
                vocab_size=self.config.vocab_size,
                **kwargs,
            )

This switch is controlled by a parameter called skip_logits, which defaults to None. If not set, it then gets evaluated to:

if skip_logits is None:
	skip_logits = self.training and (labels is not None or shift_labels is not None)

Which means that during the eval loop, the loss used will be the default one, which accumulates the logits and in turn create the huge vector mentioned above.

---

I can understand why it makes sense to accumulate logits for training algorithms like GRPO, DPO, etc., but I'm not sure it's useful for SFT unless some custom metrics are computed?

In my case, having those vRAM spikes during the eval lead to OOM errors, which defeated the purpose of using the liger-kernel in the first place.

By looking into the SFTTrainer implementation, it seemed that compute_loss method could easily be updated to force the use of the LigerForCausalLMLoss in all scenario by setting skip_logits=True:

def compute_loss(
        self,
        model: nn.Module,
        inputs: dict[str, torch.Tensor | Any],
        return_outputs: bool = False,
        num_items_in_batch: torch.Tensor | None = None,
    ):
        ...
        # Request token accuracy from Liger kernel and set token scaling if using DFT loss
        if self.args.use_liger_kernel:
            inputs["return_token_accuracy"] = True
            inputs["use_token_scaling"] = self.args.loss_type == "dft"
            
            # Force the use of LigerForCausalLMLoss instead of base ForCausalLMLoss implementation
            inputs["skip_logits"] = True

This change seems to solve the issue altogether, as shown in the vRAM graphs below (input sequence length of 40k tokens):

That's a reduction of 50GB of vRAM in our case, which is pretty substantial!

While this fix works as is, it doesn't take into account some edge cases (custom loss function provided, etc.). If you feel like this could make an interesting PR, let me know. I'd be happy to give it a crack!