SFTTrainer training very slow on GPU. Is this training speed expected?

[pledominykas]

I am currently trying to perform full fine tuning on the ai-forever/mGPT model (1.3B parameters) using a single A100 GPU (40GB VRAM) on Google Colab. However when running the training is very slow: ~0.06 it/s.

Here is my code:
`
dataset = load_dataset("allenai/c4", "lt")

train_dataset = dataset["train"]
eval_dataset = dataset["validation"]

train_dataset = train_dataset.take(10000)
eval_dataset = eval_dataset.take(1000)

trainer = SFTTrainer(
model = model,
tokenizer = tokenizer,
train_dataset = train_dataset,
eval_dataset = eval_dataset,
dataset_text_field = "text",
max_seq_length = 2048,

args = TrainingArguments(
    gradient_accumulation_steps = 4,
    gradient_checkpointing = True,

    num_train_epochs = 3,
    learning_rate = 2e-4,
    per_device_train_batch_size = 4,
    per_device_eval_batch_size = 4,

    seed = 99,
    output_dir = "./checkpoints",

    save_strategy = "steps",
    eval_strategy = "steps",

    save_steps = 0.1,
    eval_steps = 0.1,
    logging_steps = 0.1,
    load_best_model_at_end = True
),

)

trainer_stats = trainer.train()
`

And the trainer output:

It says it will take ~10hrs to process 10k examples from the c4 dataset.

These are the relevant package versions and a screenshot of GPU usage:
`
Package Version

---

accelerate 0.34.2
bitsandbytes 0.44.1
datasets 3.1.0
peft 0.13.2
torch 2.5.0+cu121
trl 0.12.0
`

It does seem to load the model to the GPU, but for some reason it’s still very slow.

I tried to use keep_in_memory=True when loading the dataset, but it did not help.

I also tried pre-tokenizing the dataset and using Trainer instead of SFTTrainer but the performance was similar.

I was wondering whether this is the expected training speed or is there some issue with my code? And if it is an issue, what could a possible fix be?

[Advait251206]

Hi @pledominykas,

Based on your configuration, I would not immediately assume that this is the expected training speed. A single A100 40GB should be capable of fine-tuning a 1.3B parameter model, although the actual throughput depends heavily on sequence length, precision, attention implementation, and data packing.

Looking at your configuration, I noticed a few things worth investigating.

1. Enable mixed precision (BF16)

Your TrainingArguments do not explicitly enable mixed precision.

If your model is currently training in FP32, this can significantly increase memory consumption and reduce throughput.

The A100 supports BF16, so I would start by enabling it:

args = TrainingArguments(
    output_dir="./checkpoints",
bf16=True,
tf32=True,

gradient_accumulation_steps=4,
gradient_checkpointing=True,

num_train_epochs=3,
learning_rate=2e-4,

per_device_train_batch_size=4,
per_device_eval_batch_size=4,

logging_steps=10,
save_strategy="steps",
save_steps=100,

eval_strategy="steps",
eval_steps=100,

seed=99,
load_best_model_at_end=True,

)

BF16 reduces the memory requirements of model computations and can improve performance on supported hardware.

tf32=True can also accelerate eligible FP32 operations on the A100.

2. Enable sequence packing

This is probably one of the first things I would investigate in your particular setup.

You are training on the C4 dataset, which contains text documents of varying lengths.

With:

max_seq_length=2048

you are setting the maximum sequence length, but that does not automatically mean your examples are packed together to utilize the full context window.

Without packing, shorter examples can leave substantial unused space in batches.

Try:

trainer = SFTTrainer(

model=model,

tokenizer=tokenizer,

train_dataset=train_dataset,
eval_dataset=eval_dataset,

dataset_text_field="text",
max_seq_length=2048,

packing=True,

args=args,

)

Packing combines multiple shorter examples into longer training sequences, potentially improving GPU utilization.

One consideration: packing changes how examples are grouped, so verify the resulting sequence lengths and attention behavior for your intended training objective.

3. Investigate gradient checkpointing

You currently have:

gradient_checkpointing=True

Gradient checkpointing reduces activation memory by recomputing certain activations during the backward pass.

However, this introduces additional computation and can slow training.

Since you have 40GB of VRAM and are already using approximately 36GB, I would first try mixed precision and packing before disabling checkpointing.

If memory permits, benchmark:

gradient_checkpointing=False

Compare the resulting throughput and peak VRAM usage.

The goal is to determine whether checkpointing is necessary for your batch size and sequence length.

4. Check the attention implementation

The attention implementation can have a significant impact on transformer training performance.

If your model architecture supports Flash Attention 2, consider loading the model with:

model = AutoModelForCausalLM.from_pretrained(

model_name,

torch_dtype=torch.bfloat16,

attn_implementation="flash_attention_2",

)

You would need to use the appropriate model class and preserve your existing tokenizer and model configuration.

Flash Attention 2 can reduce memory traffic and improve attention performance, particularly for longer sequences.

However, compatibility depends on the model architecture and installed dependencies. If mGPT does not support it, use the available attention implementation instead.

5. Check whether the GPU is actually being utilized

Your GPU memory usage is approximately 36GB, but memory allocation alone does not indicate computational utilization.

I would run the following command in another Colab cell:

nvidia-smi -l 1

Pay attention to:

• GPU utilization (%)

• GPU memory usage

• Power consumption

• GPU temperature

If GPU utilization remains low, investigate CPU preprocessing, data loading, and tokenization.

If utilization remains consistently high, the bottleneck is more likely within the training computation itself.

You can also inspect the dataset's sequence lengths after tokenization to determine whether your examples are substantially shorter than 2048 tokens.

6. Your gradient accumulation configuration

You currently have:

per_device_train_batch_size=4

gradient_accumulation_steps=4

This gives an effective batch size of 16 examples per optimizer update on a single GPU.

Gradient accumulation is useful when memory limits prevent larger physical batches, but it does not inherently improve throughput.

If memory permits, benchmark a larger per-device batch size with fewer accumulation steps while keeping the effective batch size comparable.

For example:

per_device_train_batch_size=8

gradient_accumulation_steps=2

This is only a benchmark suggestion; whether it fits depends on the actual sequence lengths and activation memory.

7. One more thing: your progress estimate

Your output shows:

[4/1875 00:39 < 10:15:06, 0.05 it/s, Epoch 0.00/3]

With 10,000 training examples, batch size 4, and gradient accumulation 4, approximately 625 optimizer updates per epoch and 1,875 over three epochs is consistent with your progress display.

However, the displayed it/s measures training iterations, not necessarily examples or tokens processed per second.

For language-model training, I would also monitor tokens per second because two configurations with identical iteration speeds can process very different amounts of text.

---

My recommended order of experimentation

Rather than changing everything simultaneously, I would benchmark these configurations individually:

Experiment	Change
Baseline	Current configuration
1	Enable bf16=True
2	Enable packing=True
3	Enable Flash Attention 2, if supported
4	Test without gradient checkpointing
5	Adjust physical batch size and accumulation

Record the training throughput, peak VRAM usage, and GPU utilization after each change.

My initial suspicion would be that the combination of unoptimized precision, lack of packing, and gradient checkpointing is contributing to the low throughput. However, the GPU utilization and actual token lengths would help identify the bottleneck more reliably.

Hope this helps!

---

If this answer helped or pointed you in the right direction, I'd appreciate it if you could mark it as the accepted answer so it's easier for others with the same issue to find.

Also, if you found my contribution useful, I'd appreciate it if you could check out my GitHub profile, follow me, and star any repositories you find interesting.

GitHub: https://github.com/Advait251206