Ways to improve my tokens / second

🤗Transformers
1

I am now pretraining a 1B model (Model: ~977M params, 24 layers, 1792 hidden, 16 heads, 112 head dim, SwiGLU, RoPE, weight-tied embeddings) from a RTX 5080 16gb VRAM, I ran it on WSL Ubuntu and also enables gradient checkpoints, BF16 autocast, 8-bit AdamW, Pytorch SDPA, and now my micro batch size is 1, gradient accumulation is 16, context is 4096; I tried to have larger micro batch but it will cause OOM; now I have a 8000-9000 tokens / second in average. Seeking for ways to make the training faster…

2

Still looking for suggestions, thank you very much.

3

Hmm… it looks like you’re already using several memory-saving and speed-up techniques, so…

Given your ~977M-parameter model, 16 GB RTX 5080, 4096-token context, microbatch size 1, and gradient accumulation of 16, I would focus on improving the efficiency of the existing training loop before changing the model or training objective.

My first two candidates would be:

  1. Try torch.compile if you are not already using it. Start with mode="default" and compare it against your current eager-mode training. This may reduce Python/kernel-launch overhead and fuse some operations, although the improvement depends on your model and implementation.
  2. Review how gradient checkpointing is applied. You already have it enabled, which saves VRAM by recomputing activations during backward. The next opportunity may be to reduce that recomputation, rather than simply enabling more memory-saving features.

I would not assume that 8,000–9,000 training tokens/s is necessarily poor for your setup. Without knowing exactly how throughput is measured and where the GPU spends its time, it is difficult to establish a meaningful target.

A practical starting point

Your nominal token count per optimizer update is:

microbatch size Ă— sequence length Ă— gradient accumulation

1 Ă— 4096 Ă— 16 = 65,536 tokens/update

If your reported 8,000–9,000 tokens/s covers all 16 forward/backward microsteps plus the optimizer update, it corresponds to approximately 7.3–8.2 seconds per optimizer update.

That gives you a useful baseline to compare against.

I would try this sequence:

  1. Confirm the steady-state time for a complete optimizer update.
  2. If torch.compile is not enabled, compare eager mode against torch.compile(model, mode="default").
  3. Record both tokens/s and peak VRAM usage.
  4. If compilation gives little improvement, examine activation checkpointing and the main GPU bottleneck.
  5. Only if the LM head/loss is consuming substantial VRAM, consider fused or chunked cross entropy to free memory for reducing checkpointing.

Keep the model, data, gradient accumulation, loss normalization, and optimizer behavior unchanged between comparisons whenever possible.

For a 16 GB GPU, I would not start with torch.compile(mode="reduce-overhead"). That mode uses CUDA Graphs and can require additional memory. The default mode is a less disruptive starting point.

References: PyTorch torch.compile, activation checkpointing.

Quick decision guide

What you observe What I would investigate
torch.compile is not being used Compare eager mode with torch.compile(..., mode="default")
Compilation is already enabled but gains are small Graph breaks, recompilations, and which parts of the model are compiled
VRAM is almost exhausted and backward recomputation is expensive Checkpointing implementation and how many layers are checkpointed
The LM head or loss creates large memory peaks Fused/chunked cross entropy, potentially followed by reduced checkpointing
GPU kernels occupy most of the time, especially attention Actual SDPA backend and attention implementation
GPU execution has long idle gaps Data loading, CPU overhead, host-to-device transfers, WSL file placement
Optimizer updates consume substantial time The specific 8-bit AdamW implementation
Your model is compatible with an optimized training framework Consider an isolated Unsloth comparison

The important distinction is that saving VRAM is not automatically the same as increasing tokens/s. Sometimes a memory-saving operation is slower on its own but enables a faster overall training configuration.

Below are the more detailed branches.

1. Measuring training throughput and locating the bottleneck

What should count as tokens/s?

For training, I would distinguish three measurements:

  • Nominal tokens/s: batch size Ă— sequence length Ă— accumulation steps, divided by the elapsed time.
  • Valid training tokens/s: tokens that actually contribute to the training loss, accounting for padding, ignored labels, and other masking.
  • End-to-end throughput: includes the relevant data loading, optimizer, scheduler, logging, and other work performed during training.

These are all useful, but they answer slightly different questions.

They should also be distinguished from inference/decode tokens per second.

For your current configuration, 65,536 is the nominal number of token positions processed per optimizer update. If some positions are ignored by the loss, that is not necessarily the number of loss-contributing tokens.

Timing a complete optimizer update

CUDA execution is asynchronous, so ordinary Python timing can underestimate elapsed GPU work unless synchronization is handled.

A minimal timing boundary could look like this:

import time
import torch

# Integration sketch:
# existing_full_update() means your current, unchanged
# accumulation-16 training update, including optimizer.step().
# It is not a replacement training-loop implementation.

torch.cuda.synchronize()
torch.cuda.reset_peak_memory_stats()

start = time.perf_counter()

existing_full_update()

torch.cuda.synchronize()
elapsed = time.perf_counter() - start

nominal_tokens = 1 * 4096 * 16

print(f"seconds/update: {elapsed:.3f}")
print(f"nominal tokens/s: {nominal_tokens / elapsed:.1f}")

print(
    "peak allocated GiB:",
    torch.cuda.max_memory_allocated() / 1024**3
)

print(
    "peak reserved GiB:",
    torch.cuda.max_memory_reserved() / 1024**3
)

The important part is the measurement boundary, not the illustrative function name.

The existing update should continue to handle your original forward passes, backward passes, gradient accumulation, optional clipping, optimizer update, scheduler, and gradient clearing in their original order.

If you use a training framework, it is preferable to instrument its existing update boundary rather than recreate the update logic.

For a reasonable comparison:

  • Exclude the initial CUDA warmup and compilation time from steady-state throughput, but account for compilation separately if total training time matters.
  • Allow the optimizer state to initialize before measuring.
  • Measure several complete updates, ideally using a median rather than a single iteration.
  • Use the same sequence lengths, data, accumulation settings, and optimizer behavior.
  • Where practical, start competing configurations from the same model and optimizer checkpoint.
  • Keep periodic evaluation and checkpoint-saving time separate from the steady training-step measurement, or include them consistently in an end-to-end comparison.

max_memory_allocated() describes the peak allocation tracked by PyTorch’s CUDA allocator. max_memory_reserved() describes memory reserved by that allocator. Neither is identical to the total GPU memory reported by nvidia-smi.

When to use a profiler

If the basic measurements do not reveal an obvious direction, a short PyTorch profiler run is more informative than trying many unrelated optimizations.

I would look for:

  • Significant gaps between GPU kernels.
  • Expensive matrix multiplications and attention operations.
  • Backward recomputation introduced by checkpointing.
  • Large LM-head projection or cross-entropy costs.
  • Expensive optimizer updates.
  • Host-to-device transfers or time spent waiting for batches.

You do not need to profile a long training run. A small number of representative optimizer updates can be sufficient for the initial distinction.

Profiling itself introduces overhead, so keep profiler timings separate from normal speed measurements.

A useful reference implementation is LitGPT’s pretraining loop, which provides examples of organizing pretraining, gradient accumulation, and performance measurement.

2. torch.compile and gradient-checkpointing options

A. If you are not using torch.compile

I would first test the simplest change:

model = torch.compile(model, mode="default")

This is an example of the intended integration point, assuming your PyTorch version, model, and training setup support compilation.

I would start by compiling the model, rather than immediately trying to compile the entire training loop, optimizer, and all 16 accumulation iterations together.

torch.compile can improve performance through graph-level optimization, operator fusion, and reductions in Python/kernel-launch overhead.

However, it does not guarantee a speedup. A 4096-token sequence still involves substantial GPU computation, and the amount of avoidable CPU overhead depends on the model implementation.

For custom architectures, the following can also matter:

  • Dynamic Python control flow.
  • Operations that cause graph breaks.
  • Changing shapes or guards that trigger recompilation.
  • Custom operators that are not handled efficiently.
  • Interactions with activation checkpointing.
  • How the optimizer and training framework are integrated.

If model-level compilation causes problems, compiling only well-behaved transformer blocks or other submodules may be a better compromise.

PyTorch’s where-to-apply-compile guide specifically discusses choosing an appropriate compilation boundary.

B. If you are already using torch.compile

Then I would not recommend adding another compilation mode blindly.

Instead, check whether compilation is actually covering the expensive parts of your model and whether it repeatedly recompiles.

For a short diagnostic run, PyTorch’s compilation logs can help:

TORCH_LOGS="graph_breaks,recompiles" python train.py

The relevant logging options depend on the installed PyTorch version.

Useful references:

If compilation is already effective and your GPU is heavily occupied with useful computation, additional launch-overhead optimization may not be the highest-value direction.

C. Checkpointing is already enabled — optimize how it is used

Since you already use gradient checkpointing, the interesting question is how much recomputation it introduces and whether some of it can be avoided.

There are two related but separate experiments.

First: check the checkpoint implementation.

PyTorch provides reentrant and non-reentrant checkpointing through use_reentrant.

Its documentation recommends use_reentrant=False for the general case.

The non-reentrant implementation can stop recomputing once the necessary intermediate activations have been recovered, whereas the reentrant implementation recomputes the checkpointed function in full.

If you currently use reentrant checkpointing, comparing the non-reentrant path may be worthwhile.

If you already use non-reentrant checkpointing, there is no reason to repeat the same change.

This is not a universal speed guarantee. Behavior depends on the checkpointed function, the autograd graph, the installed PyTorch version, and interactions with compilation.

See PyTorch’s checkpointing documentation.

Second: experiment with checkpointing coverage.

You have 24 transformer layers.

If every layer is checkpointed, and there is enough memory headroom, you could test checkpointing slightly fewer layers rather than disabling checkpointing entirely.

The general trade-off is:

More checkpointing
    -> lower activation memory
    -> more backward recomputation

Less checkpointing
    -> higher activation memory
    -> potentially less backward recomputation

With only 16 GB VRAM and microbatch sizes above 1 already causing OOM, I would not begin by disabling all checkpointing.

A limited, reversible adjustment is more informative and less disruptive.

One change at a time is preferable: compare the checkpointing implementation separately from changing its coverage.

If the measurements show substantial recomputation overhead, selective activation checkpointing is another, more advanced option. It can allow different classes of operations to receive different save/recompute treatment.

That would be a later optimization rather than my first suggestion.

D. What about CUDA Graphs?

torch.compile(mode="reduce-overhead") can reduce CPU launch overhead through CUDA Graphs.

However, the documentation notes that it may cache additional workspace memory, and CUDA Graph capture has compatibility constraints.

On a GPU where memory is already tight, I would try this only if measurements suggest that launch overhead is important and there is enough memory headroom.

The first comparison should remain eager versus mode="default".

3. If the LM head and cross-entropy loss are taking too much memory

This is a separate optimization path that might be especially interesting if your vocabulary is large.

For a causal LM, the output projection commonly produces logits with a shape equivalent to:

[batch_size Ă— sequence_length, vocabulary_size]

With microbatch 1 and sequence length 4096, the logits tensor can become substantial.

For illustration, a single BF16 logits tensor would occupy approximately:

Assumed vocabulary Raw BF16 logits
32,768 tokens 256 MiB
65,536 tokens 512 MiB

These are illustrative vocabulary sizes, not assumptions about your model.

They describe only one logits tensor. Additional intermediate tensors, gradients, and temporary buffers can increase the peak memory footprint.

Option A: Liger Kernel’s fused linear cross entropy

Liger Kernel includes fused linear cross-entropy implementations designed to avoid materializing the full logits tensor in the conventional way.

One useful property is that its fused loss can be integrated individually, without necessarily replacing the entire model or training framework.

See the fused linear cross-entropy implementation.

This could be useful if the LM head is contributing significantly to your peak VRAM usage.

There is an important qualification, though: fused loss does not automatically mean faster training.

For reference, in a separate small synthetic experiment on an NVIDIA L4, I compared ordinary PyTorch linear projection plus cross entropy against Liger’s fused linear cross entropy.

The experiment used 4096 token positions, hidden size 1792, an assumed vocabulary of 32,768, and BF16.

Head-only measurement Ordinary PyTorch CE Liger fused CE
Median forward + backward time 29.3 ms 40.0 ms
Incremental peak allocated memory 768 MiB 142 MiB

These were eight measured iterations per implementation after warmup, using PyTorch 2.11.0 and Liger 0.8.4.

The fused implementation significantly reduced the memory required by the tested head/loss operation, but the operation itself was slower.

With a larger synthetic vocabulary of 65,536, the memory reduction was greater, while the head-only execution-time cost also increased.

This experiment did not include the 24 transformer layers, gradient accumulation of 16, the optimizer, or your RTX 5080. It therefore cannot predict your full training throughput.

There is also a separate Liger issue #323 describing situations where memory improvements did not translate into faster training. It is a different environment and should not be treated as evidence of the same cause in your setup.

The more interesting possibility: turn saved memory into speed

Suppose a fused loss saves enough VRAM to allow checkpointing fewer transformer layers.

Then the overall trade-off becomes:

Ordinary head/loss
    + more activation checkpointing
    + more backward recomputation

versus

Memory-efficient head/loss
    + less activation checkpointing
    + less backward recomputation

The fused head may be slower on its own, yet the complete optimizer update could become faster if the reduction in recomputation exceeds the extra head cost.

Or it could become slower overall.

That is why I would compare complete optimizer-update time, not just the loss kernel’s latency.

If switching to a memory-efficient loss does not allow you to reduce checkpointing or another costly memory-saving operation, the additional implementation complexity may not be worthwhile.

Option B: PyTorch’s chunked linear cross entropy

Recent PyTorch versions also document torch.nn.functional.linear_cross_entropy, with chunking configured through LinearCrossEntropyOptions.

This is another possible way to reduce output-layer memory.

It is worth checking your installed PyTorch version before considering this path, because the API is not available in every older release.

A few implementation details matter:

  • Chunking requires the appropriate non-None options configuration; the default reference path is not the same optimization.
  • Sequence representations may need to be flattened to the expected two-dimensional input.
  • The documented chunked path has limitations involving higher-order/forward-mode automatic differentiation and some torch.func operations.
  • Current documentation indicates that torch.compile falls back to eager execution at the chunked operation.

Those constraints do not make the technique unusable for ordinary training, but they mean its interaction with an otherwise compiled model should be measured.

See the API documentation and chunking options.

Preserve the training behavior when replacing the loss

Since your model uses weight-tied embeddings, I would be particularly careful about the gradient contract when replacing the output head or loss.

In a tied-weight model, the same parameter receives contributions through both the input embedding path and the output projection path.

A replacement should preserve that relationship.

Other things worth checking are:

  • The next-token target shift: make sure it is applied exactly once.
  • ignore_index and attention/loss masking semantics.
  • Whether the loss uses mean, sum, or a custom normalization.
  • How the loss is scaled over 16 accumulation microsteps.
  • Behavior when different microsteps contain different numbers of valid labels.
  • Gradient contributions to the tied parameter.
  • BF16 numerical differences and finite gradients.

In particular, averaging 16 independently normalized microbatch losses is not necessarily the same as globally averaging over all valid tokens if their valid-token counts differ.

That does not mean your current approach is wrong. It means an optimization should preserve your existing intended update rule rather than silently introduce a different one.

A small FP32 loss/gradient comparison before a BF16 full-training A/B test is a useful safeguard.

4. Other possible bottlenecks: attention, CPU/data pipeline, optimizer, and WSL

These would be conditional branches, not things I would change simultaneously.

A. SDPA is enabled, but which attention kernel is actually running?

Using torch.nn.functional.scaled_dot_product_attention is already a useful step.

However, PyTorch can select different SDPA implementations depending on the device, dtype, tensor shapes, masks, and other conditions.

Therefore, calling SDPA does not by itself prove which optimized backend executed.

See PyTorch’s SDPA documentation.

If a profiler shows that attention dominates the update, I would check the actual kernel names and relevant attention configuration.

For example, the causal masking arrangement, dropout behavior, tensor layouts, and supported fused-attention paths can influence the implementation selected.

I would not change masking semantics merely to select a faster kernel, and I would not add a separate FlashAttention package before confirming that attention is actually a bottleneck.

If your attention path is already efficient, another part of the model may deserve more attention.

B. If GPU idle time is substantial

Then the problem may not be GPU arithmetic at all.

Possible areas to inspect include:

  • On-the-fly tokenization or text processing.
  • DataLoader worker behavior and batch preparation.
  • Host-to-device copies.
  • CPU scheduling and Python overhead.
  • Dataset storage and filesystem access.
  • Whether data is being read from the WSL Linux filesystem or through a Windows-mounted path.

Since you mentioned WSL Ubuntu, file placement is worth considering if measurements reveal significant input or filesystem waiting.

However, WSL itself should not be assumed to be the cause.

NVIDIA documents several CUDA-on-WSL limitations, including aspects of pinned memory and Unified Memory behavior:

NVIDIA CUDA on WSL guide.

A small experiment using preloaded or already prepared batches can help distinguish data supply from GPU computation, without redesigning the entire pipeline.

If the GPU remains busy with useful kernels and data transfers are not causing delays, this branch becomes much less important.

C. 8-bit AdamW saves memory, but optimizer time is a separate question

Your 8-bit AdamW setup is relevant to the 16 GB memory constraint.

However, optimizer-state compression and optimizer execution speed are distinct properties.

Some small parameter tensors may also retain higher-precision optimizer states, depending on the implementation.

For example, bitsandbytes documents a default threshold below which parameter tensors remain in 32-bit optimizer state:

bitsandbytes 8-bit optimizer documentation.

If profiling shows that optimizer updates take a significant fraction of the total time, comparing implementations may help.

I would keep the same optimizer algorithm, hyperparameters, and intended update behavior during that comparison, rather than combining it with a switch to a different optimization algorithm.

If optimizer time is small relative to the 16 forward/backward microsteps, this is probably not the first place to spend effort.

D. BF16 autocast does not necessarily mean BF16 parameter storage

One additional memory detail: using BF16 autocast does not automatically imply that the model parameters are stored in BF16.

For example:

print(next(model.parameters()).dtype)

The parameter-storage dtype, compute dtype, gradient dtype, and optimizer-state representation can differ.

That distinction can affect a nearly full 16 GB VRAM budget.

However, changing stored FP32 parameters to BF16 is not merely a performance-neutral implementation change. It may change numerical behavior and the training trajectory.

I would inspect the current dtype before considering such a change, and keep it separate from the initial torch.compile and checkpointing comparisons.

5. Where Unsloth might fit

Unsloth is another reasonable implementation option to investigate.

Its documentation covers optimized training paths, including full fine-tuning and continued pretraining:

There is one distinction I would make before recommending a migration.

Pretraining a custom architecture from random initialization is not necessarily the same workflow as continued pretraining of a supported pretrained model.

Your description provides the architecture dimensions, but not the model class or initialization path.

If you are using a supported Hugging Face model architecture and training route, Unsloth may be worth testing as a separate implementation.

If this is a custom nn.Module trained from scratch, I would first confirm that Unsloth’s relevant kernels and training path actually support that model and the intended full-parameter training behavior.

The presence of pretraining examples in the documentation does not by itself establish compatibility with every custom architecture.

I would also keep these concepts separate:

  • Full-parameter training.
  • LoRA/QLoRA fine-tuning.
  • Continued pretraining from existing model weights.
  • Training from random initialization.
  • Quantized model-weight loading.
  • An 8-bit optimizer for training model parameters.

They solve overlapping but different problems.

A faster LoRA or continued-pretraining example is not necessarily an appropriate throughput prediction for full-parameter pretraining of your model.

If Unsloth supports your exact architecture and training objective, it would make sense to compare its complete optimizer-update throughput and memory use against the existing implementation.

Otherwise, retaining your current training loop while selectively using optimized PyTorch or fused kernels may be simpler and easier to validate.

What I would prioritize overall

For your particular constraints, my order would be:

First: establish a consistent full-update throughput measurement.

Second: compare eager mode with torch.compile(mode="default"), unless you have already done so.

Third: review checkpointing implementation and coverage, especially if memory pressure forces substantial recomputation.

Fourth: if the LM head is a significant memory contributor, test whether a memory-efficient loss creates enough headroom to reduce checkpointing and improve overall throughput.

After that: investigate attention, data loading, optimizer overhead, or a compatible optimized framework according to what the measurements actually show.

I would evaluate changes using three outcomes together: steady training tokens/s, peak VRAM, and preservation of the intended loss/gradient behavior. A change that reduces memory but slows the complete update may still be valuable if it unlocks a faster configuration; otherwise it may not be worth keeping.

Two details would help narrow down the most useful branch:

  • Are you already using torch.compile, and if so, in which mode?
  • Does your 8,000–9,000 tokens/s figure include all 16 accumulation microsteps and the optimizer update, measured after warmup?

Even without those details, the comparisons above can be tried incrementally without changing the training objective or replacing the whole pipeline.

4

Thank you for such a comprehensive and detailed breakdown! I really appreciate the deep dive.

I just finished running some A/B tests based on your suggestions, and here are the actual results on my setup (RTX 5080 16GB, PyTorch 2.14, CUDA 13):

  1. torch.compile: I tested eager vs. mode="default". In my WSL environment with custom checkpointing, it caused graph breaks and was actually slightly slower than eager mode.

  2. Gradient Checkpointing: I tried reducing the checkpointed layers, but since my baseline VRAM is already at 15.0/16.0 GB, any reduction immediately caused an OOM (Out of Memory). So I had to keep all 24 layers checkpointed.

  3. Fused Loss (Liger): I tried to implement Liger’s fused linear cross entropy. Unfortunately, PyTorch 2.14 + CUDA 13 seems to have a deep compatibility issue with Liger. It threw CUBLAS_STATUS_INTERNAL_ERROR and illegal memory access errors, even in the nightly build.

Current Status: I reverted to the native F.cross_entropy and it’s running perfectly stable at ~9,000 tokens/s with peak VRAM.

Thank you again for the thorough guide! It helped me rule out all the dead ends and find the most stable configuration for my hardware.